57 lines
2.3 KiB
Python
57 lines
2.3 KiB
Python
import argparse
|
||
import os
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
from segment_anything import sam_model_registry, SamPredictor
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="SAM box prompt segmentation 批量处理")
|
||
parser.add_argument("--input-dir", type=str, required=True, help="输入图片文件夹路径")
|
||
parser.add_argument("--box-file", type=str, required=True, help="box信息文件路径，每行: 图片名 x1 y1 x2 y2")
|
||
parser.add_argument("--checkpoint", type=str, required=True, help="SAM模型权重路径")
|
||
parser.add_argument("--model-type", type=str, default="vit_h", help="模型类型，默认vit_h")
|
||
parser.add_argument("--output-dir", type=str, required=True, help="输出掩码图片文件夹路径")
|
||
args = parser.parse_args()
|
||
|
||
# 检查CUDA
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
if device == "cpu":
|
||
print("警告：未检测到CUDA，正在使用CPU，速度会较慢。")
|
||
|
||
# 加载模型
|
||
sam = sam_model_registry[args.model_type](checkpoint=args.checkpoint)
|
||
sam.to(device)
|
||
predictor = SamPredictor(sam)
|
||
|
||
# 创建输出文件夹
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
||
# 读取box文件并批量处理
|
||
with open(args.box_file, 'r', encoding='utf-8') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
parts = line.split()
|
||
if len(parts) != 5:
|
||
print(f"格式错误: {line}")
|
||
continue
|
||
img_name, x1, y1, x2, y2 = parts
|
||
box = np.array([int(x1), int(y1), int(x2), int(y2)])
|
||
img_path = os.path.join(args.input_dir, img_name)
|
||
image = cv2.imread(img_path)
|
||
if image is None:
|
||
print(f"无法读取图片: {img_path}")
|
||
continue
|
||
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
predictor.set_image(image_rgb)
|
||
masks, scores, logits = predictor.predict(box=box[None, :], multimask_output=False)
|
||
mask = masks[0].astype(np.uint8) * 255
|
||
out_path = os.path.join(args.output_dir, f"{os.path.splitext(img_name)[0]}_mask.png")
|
||
cv2.imwrite(out_path, mask)
|
||
print(f"掩码已保存到: {out_path}")
|
||
|
||
if __name__ == "__main__":
|
||
main() |