minsung
4e5173522d
feat: SAM3.1 다중 프롬프트 배치 검출 및 라벨 병합 후처리
sam3_multi_prompt.py: 프롬프트 N개를 forward 1회에 배치 처리한다.
서버(프롬프트당 forward 1회) 대비 forward 횟수가 1/N로 줄고, 이미지
임베딩과 텍스트 인코딩을 타일 전체에서 재사용한다. RTX 3060 12GB
기준 배치 16이 최적(32 이상은 VRAM 압박으로 4배 이상 느려짐).
타일보다 큰 대상(숲, 도로, 주차장)은 타일 경계에서 잘려 사각형
마스크가 되므로 --wide-prompts 로 지정해 이미지 전체를 한 장으로
처리한다.
merge_labels.py: 검출 결과 후처리. 같은 병합 그룹에 속한 라벨끼리
외곽선이 gap px 이내로 인접하면 하나로 합치고 대표 라벨을 붙인다
(building + building rooftop + blue roof -> building). 병합 전
claim 규칙으로 소유권을 재배정해, 차량 위에 잡힌 *roof 폴리곤이
건물이 아니라 차량에 합쳐지도록 한다.
make_viewer.py: 라벨별 on/off 가능한 단독 HTML 뷰어 생성.
프롬프트 파일의 그룹 주석 기준으로 라벨을 묶어 나열하고, 줌/팬과
선택 상태 저장을 지원한다.
cut_tiles.py: R{행}C{열} 격자 오버레이 + 원본 배율 타일 저장.
전역 규칙 반영:
- CUDA 사용 불가 시 CPU로 폴백하지 않고 즉시 에러
- set_per_process_memory_fraction 으로 드라이버의 시스템 메모리
폴백(10배 이상 느려짐) 전에 OOM 발생
- 체크포인트/BPE 경로 하드코딩 제거. CLI 인자 -> 환경변수 ->
sam31server 설정파일 순으로 해결
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@
299 lines
13 KiB
Python
299 lines
13 KiB
Python
"""
|
||
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()
|