""" SAM3.1 다중 프롬프트 배치 세그멘테이션 (in-process, 서버 불필요) 프롬프트 N개를 forward 1회에 함께 처리한다. 서버 방식(프롬프트당 forward 1회)과 달리 이미지 임베딩·텍스트 인코딩을 재사용하므로 프롬프트 수가 많을수록 유리하다. 사용법: D:/MYCLAUDE_PROJECT/sam31server/.venv/Scripts/python.exe tools/sam3_multi_prompt.py \ --input "data/everyimage/xxx.JPG" \ --prompts prompts/discovery_v1.txt \ --cols 9 --rows 6 --conf 0.25 --merge 사전 조건: SAM3 서버는 내려둘 것 (GPU에 모델 2벌 올라감) """ import argparse import json import os import sys import time from collections import Counter from pathlib import Path import cv2 import numpy as np import torch from PIL import Image SERVER_PATH = Path(os.environ.get( "SAM31SERVER_DIR", Path(__file__).resolve().parent.parent.parent / "sam31server")) if not SERVER_PATH.is_dir(): raise SystemExit(f"sam31server 없음: {SERVER_PATH} (SAM31SERVER_DIR 로 지정)") sys.path.insert(0, str(SERVER_PATH)) sys.path.insert(0, str(Path(__file__).resolve().parent)) from sam3_everything_explore import ( # noqa: E402 후처리·시각화 재사용 _bbox, nms_shapes, merge_adjacent, draw_everything, analyze_labels, ) def resolve_model_paths(server_dir: Path, ckpt_arg=None, bpe_arg=None): """체크포인트·BPE 경로 해결: CLI 인자 → 환경변수 → sam31server 설정파일.""" ckpt, bpe = ckpt_arg or os.environ.get("SAM3_CHECKPOINT"), \ bpe_arg or os.environ.get("SAM3_BPE") if not (ckpt and bpe): cfg = server_dir / "configs" / "auto_labeling" / "segment_anything_3.yaml" if not cfg.is_file(): raise SystemExit(f"설정 파일 없음: {cfg} (--checkpoint / --bpe 로 직접 지정)") import yaml params = yaml.safe_load(cfg.read_text(encoding="utf-8")).get("params", {}) ckpt = ckpt or params.get("model_path") bpe = bpe or params.get("bpe_path") for name, p in (("체크포인트", ckpt), ("BPE 사전", bpe)): if not p or not Path(p).is_file(): raise SystemExit(f"{name} 파일 없음: {p}") return ckpt, bpe def load_prompts(path: Path) -> list: """# 주석과 빈 줄을 제외한 프롬프트 목록.""" lines = path.read_text(encoding="utf-8").splitlines() return [ln.strip() for ln in lines if ln.strip() and not ln.strip().startswith("#")] def tile_boxes(W, H, cols, rows, overlap): """(x0, y0, x1, y1) 타일 목록. overlap 비율만큼 확장.""" bw, bh = W / cols, H / rows px, py = int(bw * overlap), int(bh * overlap) boxes = [] for r in range(rows): for c in range(cols): boxes.append(( max(0, int(c * bw) - px), max(0, int(r * bh) - py), min(W, int((c + 1) * bw) + px), min(H, int((r + 1) * bh) + py), )) return boxes def masks_to_polygons(masks, epsilon_factor=0.001): """[K,h,w] bool 텐서 → 폴리곤 리스트 (없으면 None).""" polys = [] for m in masks: mu = m.astype(np.uint8) contours, _ = cv2.findContours(mu, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: polys.append(None) continue largest = max(contours, key=cv2.contourArea) eps = epsilon_factor * cv2.arcLength(largest, True) approx = cv2.approxPolyDP(largest, eps, True) polys.append(approx if len(approx) >= 3 else None) return polys def predict_tile(model, processor, find_stage_cls, tile_bgr, text_outs, chunks, conf, mask_bytes=6 * 10**8): """타일 1장에 프롬프트 전체를 배치로 물어본다. 반환: shape dict 리스트.""" th, tw = tile_bgr.shape[:2] # 업샘플 한 번에 올릴 마스크 수 — 타일이 클수록 줄인다 (통짜 패스 OOM 방지) mask_batch = max(1, min(32, mask_bytes // (th * tw * 4))) state = processor.set_image(Image.fromarray(tile_bgr[:, :, ::-1])) shapes = [] for captions, text_out in zip(chunks, text_outs): n = len(captions) state["backbone_out"].update(text_out) find = find_stage_cls( img_ids=torch.zeros(n, dtype=torch.long, device=model.device), text_ids=torch.arange(n, dtype=torch.long, device=model.device), input_boxes=None, input_boxes_mask=None, input_boxes_label=None, input_points=None, input_points_mask=None, ) out = model.forward_grounding( backbone_out=state["backbone_out"], find_input=find, geometric_prompt=model._get_dummy_prompt(num_prompts=n), find_target=None, ) probs = out["pred_logits"].sigmoid() # [n,Q,1] presence = out["presence_logit_dec"].sigmoid().unsqueeze(1) probs = (probs * presence).squeeze(-1) # [n,Q] keep = probs > conf idx = keep.nonzero(as_tuple=False) if idx.numel() == 0: continue sel_masks = out["pred_masks"][keep] # [K,mh,mw] sel_scores = probs[keep] # 마스크를 타일 크기로 키운 뒤 외곽선을 뽑는다. 원본 해상도(252px 정도)에서 # 뽑으면 좌표가 격자에 박혀 계단 현상이 생긴다. 메모리 때문에 조각내서 처리. for s in range(0, sel_masks.shape[0], mask_batch): chunk = sel_masks[s:s + mask_batch].unsqueeze(1).float() up = torch.nn.functional.interpolate( chunk, (th, tw), mode="bilinear", align_corners=False) binary = (up > 0).squeeze(1).cpu().numpy().astype(np.uint8) for k, poly in enumerate(masks_to_polygons(binary)): if poly is None: continue b = int(idx[s + k, 0]) shapes.append({ "label": captions[b], "score": float(sel_scores[s + k]), "shape_type": "polygon", "points": [[float(p[0][0]), float(p[0][1])] for p in poly], }) del chunk, up del out return shapes def run_pass(model, processor, find_stage_cls, image_bgr, boxes, captions, conf, batch, tag, mask_bytes=6 * 10**8): """타일 목록 전체에 프롬프트 집합을 돌린다. 반환: 전역 좌표 shape 리스트.""" chunks = [captions[i:i + batch] for i in range(0, len(captions), batch)] shapes, t0 = [], time.time() with torch.inference_mode(): text_outs = [model.backbone.forward_text(c, device=model.device) for c in chunks] for i, (x0, y0, x1, y1) in enumerate(boxes, 1): got = predict_tile(model, processor, find_stage_cls, image_bgr[y0:y1, x0:x1], text_outs, chunks, conf, mask_bytes) for s in got: # 전역 좌표로 이동 s["points"] = [[p[0] + x0, p[1] + y0] for p in s["points"]] shapes.extend(got) torch.cuda.empty_cache() # 타일 간 VRAM 누적 방지 print(f" [{tag}] 타일 {i}/{len(boxes)} +{len(got)}개 " f"(누적 {len(shapes)}, {time.time()-t0:.0f}초, " f"VRAM {torch.cuda.memory_reserved()/2**30:.1f}GB)") return shapes def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--input", default=None) ap.add_argument("--output", default=None, help="단일 이미지용 출력 경로") ap.add_argument("--outdir", default=None, help="결과 폴더. --input 이 폴더면 필수 — <이름>_multi.jpg/.json 로 저장") ap.add_argument("--prompts", default="prompts/discovery_v1.txt") ap.add_argument("--wide-prompts", default=None, help="타일보다 큰 대상 목록. 여기 적힌 라벨은 타일 패스에서 빼고 " "이미지 전체를 한 장으로 검출한다 (예: prompts/wide_v1.txt)") ap.add_argument("--wide-in-tiles", action="store_true", help="통짜 프롬프트를 타일 패스에서도 돌린다. 통짜는 대상 하나에 " "인스턴스가 하나만 살아남아 도로 한쪽 차로·논 한 필지만 " "잡히는 일이 있다. 겹치는 결과는 NMS 가 지운다. 느려진다") ap.add_argument("--cols", type=int, default=9) ap.add_argument("--rows", type=int, default=6) ap.add_argument("--overlap", type=float, default=0.10) ap.add_argument("--conf", type=float, default=0.25) ap.add_argument("--nms", type=float, default=0.40) ap.add_argument("--device", default="cuda:0", help="사용할 GPU (예: cuda:0, cuda:1). --list-gpus 로 확인") ap.add_argument("--list-gpus", action="store_true", help="GPU 목록만 출력하고 종료") ap.add_argument("--batch", type=int, default=0, help="forward 1회에 넣을 최대 프롬프트 수. 0이면 VRAM에서 자동 결정 " "(12GB→16, 24GB→32). 3060 12GB에서 32는 4배 이상 느려짐") ap.add_argument("--merge", action="store_true", help="같은 라벨 인접 폴리곤 병합") ap.add_argument("--merge-gap", type=int, default=8) ap.add_argument("--checkpoint", default=None, help="SAM3.1 체크포인트 (기본: 환경변수 SAM3_CHECKPOINT → 서버 설정파일)") ap.add_argument("--bpe", default=None, help="BPE 사전 (기본: 환경변수 SAM3_BPE → 서버 설정파일)") ap.add_argument("--vram-fraction", type=float, default=0.92, help="VRAM 사용 상한 비율. 넘으면 느려지는 대신 OOM 에러 (기본 0.92)") args = ap.parse_args() if not torch.cuda.is_available(): raise SystemExit("CUDA 사용 불가. CPU 폴백하지 않는다 — GPU 환경을 확인하라.") if args.list_gpus: for i in range(torch.cuda.device_count()): p = torch.cuda.get_device_properties(i) print(f" cuda:{i} {p.name} {p.total_memory/2**30:.0f}GB") return if not args.input: raise SystemExit("--input 필요") idx = torch.device(args.device).index or 0 if idx >= torch.cuda.device_count(): raise SystemExit(f"{args.device} 없음 — 장착 GPU {torch.cuda.device_count()}개. " f"--list-gpus 로 확인하라.") torch.cuda.set_device(idx) # model_builder 가 device == "cuda" 문자열을 정확 비교하므로 "cuda:N" 을 주면 # 모델이 CPU에 남는다. 장치 선택은 set_device 로 하고 문자열은 "cuda" 고정. device = "cuda" vram_gb = torch.cuda.get_device_properties(idx).total_memory / 2**30 if args.batch <= 0: # VRAM에 맞춰 자동 결정 args.batch = 32 if vram_gb >= 20 else 16 print(f"GPU cuda:{idx} {torch.cuda.get_device_name(idx)} " f"{vram_gb:.0f}GB batch={args.batch}") captions = load_prompts(Path(args.prompts)) if not captions: print(f"프롬프트 없음: {args.prompts}") return src = Path(args.input) if src.is_dir(): images = sorted(p for p in src.iterdir() if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".tif")) if not images: raise SystemExit(f"이미지 없음: {src}") if not args.outdir: raise SystemExit("폴더 입력에는 --outdir 필요") else: images = [src] wide = load_prompts(Path(args.wide_prompts)) if args.wide_prompts else [] fine = captions if args.wide_in_tiles else [c for c in captions if c not in set(wide)] def nchunk(n): return (n + args.batch - 1) // args.batch print(f"입력 : {len(images)}장") print(f"타일 : {args.cols}×{args.rows}={args.cols*args.rows}개 " f"overlap={args.overlap*100:.0f}%") print(f"타일 패스 : 프롬프트 {len(fine)}개 → forward " f"{nchunk(len(fine))*args.cols*args.rows}회/장") if wide: print(f"통짜 패스 : 프롬프트 {len(wide)}개 → forward {nchunk(len(wide))}회/장") print(f"conf={args.conf} nms={args.nms}\n") from sam3.model_builder import build_sam3_image_model from sam3.model.sam3_image_processor import Sam3Processor from sam3.model.data_misc import FindStage # VRAM 상한을 걸어 드라이버가 시스템 메모리로 폴백(10배 이상 느려짐)하기 전에 # OOM으로 실패하게 만든다 torch.cuda.set_per_process_memory_fraction(args.vram_fraction, idx) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True ckpt_path, bpe_path = resolve_model_paths(SERVER_PATH, args.checkpoint, args.bpe) print(f"체크포인트: {ckpt_path}") print("SAM3.1 로딩...") model = build_sam3_image_model( bpe_path=bpe_path, device=device, checkpoint_path=ckpt_path) processor = Sam3Processor(model, confidence_threshold=args.conf, device=device) mask_bytes = int(6 * 10**8 * vram_gb / 12) # VRAM에 비례해 업샘플 청크 조정 t_all = time.time() for n, img_path in enumerate(images, 1): print(f"\n===== [{n}/{len(images)}] {img_path.name} =====") image_bgr = cv2.imdecode(np.fromfile(str(img_path), dtype=np.uint8), cv2.IMREAD_COLOR) if image_bgr is None: raise SystemExit(f"이미지 로드 실패: {img_path}") H, W = image_bgr.shape[:2] boxes = tile_boxes(W, H, args.cols, args.rows, args.overlap) t0 = time.time() all_shapes = run_pass(model, processor, FindStage, image_bgr, boxes, fine, args.conf, args.batch, "타일", mask_bytes) if wide: # 타일보다 큰 대상은 이미지 전체를 한 장으로 보고 검출 all_shapes += run_pass(model, processor, FindStage, image_bgr, [(0, 0, W, H)], wide, args.conf, args.batch, "통짜", mask_bytes) print(f" 검출 {len(all_shapes)}개 → NMS(iou={args.nms})...") all_shapes = nms_shapes(all_shapes, iou_thresh=args.nms) print(f" NMS 후 {len(all_shapes)}개") if args.merge: before = len(all_shapes) all_shapes = merge_adjacent(all_shapes, gap=args.merge_gap) print(f" 병합(gap={args.merge_gap}px) {before} → {len(all_shapes)}개") print(f" {time.time()-t0:.0f}초") vis = draw_everything(image_bgr, all_shapes, args.cols, args.rows) h, w = vis.shape[:2] if max(h, w) > 4096: s = 4096 / max(h, w) vis = cv2.resize(vis, (int(w * s), int(h * s))) if args.outdir: out_path = Path(args.outdir) / (img_path.stem + "_multi.jpg") elif args.output: out_path = Path(args.output) else: out_path = img_path.parent / (img_path.stem + "_multi.jpg") out_path.parent.mkdir(parents=True, exist_ok=True) cv2.imencode(".jpg", vis, [cv2.IMWRITE_JPEG_QUALITY, 93])[1].tofile(str(out_path)) json_path = out_path.with_suffix(".json") json_path.write_text(json.dumps({ "source_image": str(img_path), "total_segments": len(all_shapes), "label_counts": dict(Counter(s.get("label", "") for s in all_shapes)), "segments": [{"label": s.get("label", ""), "score": s.get("score", 0), "bbox": list(_bbox(s["points"])), "points": s["points"]} for s in all_shapes], }, ensure_ascii=False, indent=2), encoding="utf-8") print(f" 저장: {out_path.name} / {json_path.name}") print(f"\n전체 {len(images)}장 완료 — {time.time()-t_all:.0f}초") if sys.platform == "win32": # 완료 알림음 import winsound winsound.Beep(880, 150) winsound.Beep(1175, 250) if __name__ == "__main__": main()