""" 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): """타일 목록 전체에 프롬프트 집합을 돌린다. 반환: 전역 좌표 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) 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", required=True) ap.add_argument("--output", default=None, help="기본: 입력명_multi.jpg") ap.add_argument("--prompts", default="prompts/discovery_v1.txt") ap.add_argument("--wide-prompts", default=None, help="타일보다 큰 대상 목록. 여기 적힌 라벨은 타일 패스에서 빼고 " "이미지 전체를 한 장으로 검출한다 (예: prompts/wide_v1.txt)") 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("--batch", type=int, default=16, help="forward 1회에 넣을 최대 프롬프트 수 (기본 16). " "RTX 3060 12GB 기준 16이 최적 — 32 이상은 VRAM 압박으로 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() captions = load_prompts(Path(args.prompts)) if not captions: print(f"프롬프트 없음: {args.prompts}") return img_path = Path(args.input) image_bgr = cv2.imdecode(np.fromfile(str(img_path), dtype=np.uint8), cv2.IMREAD_COLOR) if image_bgr is None: print(f"이미지 로드 실패: {img_path}") return H, W = image_bgr.shape[:2] wide = load_prompts(Path(args.wide_prompts)) if args.wide_prompts else [] fine = [c for c in captions if c not in set(wide)] boxes = tile_boxes(W, H, args.cols, args.rows, args.overlap) def nchunk(n): return (n + args.batch - 1) // args.batch print(f"이미지 : {W}×{H}") print(f"타일 : {args.cols}×{args.rows}={len(boxes)}개 overlap={args.overlap*100:.0f}%") print(f"타일 패스 : 프롬프트 {len(fine)}개 → forward {nchunk(len(fine))*len(boxes)}회") if wide: print(f"통짜 패스 : 프롬프트 {len(wide)}개 → forward {nchunk(len(wide))}회 " f"(타일보다 큰 대상)") 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 if not torch.cuda.is_available(): raise SystemExit("CUDA 사용 불가. CPU 폴백하지 않는다 — GPU 환경을 확인하라.") device = "cuda" # VRAM 상한을 걸어 드라이버가 시스템 메모리로 폴백(10배 이상 느려짐)하기 전에 # OOM으로 실패하게 만든다 torch.cuda.set_per_process_memory_fraction(args.vram_fraction) 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) t0 = time.time() all_shapes = run_pass(model, processor, FindStage, image_bgr, boxes, fine, args.conf, args.batch, "타일") if wide: # 타일보다 큰 대상은 이미지 전체를 한 장으로 보고 검출 all_shapes += run_pass(model, processor, FindStage, image_bgr, [(0, 0, W, H)], wide, args.conf, args.batch, "통짜") print(f"\n검출 {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}초\n") analyze_labels(all_shapes) 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))) out_path = (Path(args.output) if args.output else 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)) print(f"\n저장: {out_path}") json_path = out_path.with_suffix(".json") json_path.write_text(json.dumps({ "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"라벨 데이터: {json_path}") if sys.platform == "win32": # 완료 알림음 import winsound winsound.Beep(880, 150) winsound.Beep(1175, 250) if __name__ == "__main__": main()