diff --git a/.gitignore b/.gitignore index 5dfb8f8..52912b8 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,9 @@ data/ # 논문/문서 파일 *.pdf *.txt +# 단, 프롬프트·병합 설정은 코드의 일부라 포함 +!prompts/*.txt +!configs/*.txt # 미정리 툴 (작업 중) tools/render_polygons_rainbow.py diff --git a/configs/claim_rules.txt b/configs/claim_rules.txt new file mode 100644 index 0000000..7ef9f51 --- /dev/null +++ b/configs/claim_rules.txt @@ -0,0 +1,7 @@ +# 소유권 재배정 규칙 — 병합 전에 실행된다. +# 형식: <가져갈 그룹> : <최소 겹침 비율> : <대상 라벨, 콤마 구분> +# +# 대상 라벨의 폴리곤이 그 그룹 폴리곤과 자기 면적의 <비율> 이상 겹치면, +# 겹친 상대의 라벨로 바꿔 붙인다. (예: 자동차 위의 "white roof" → "white vehicle") + +vehicle : 0.5 : blue roof, red roof, green roof, gray roof, dark gray roof, black roof, white roof, silver metal roof, orange roof, brown roof, yellow roof, blue waterproof tarp, black protective sheet diff --git a/configs/merge_groups.txt b/configs/merge_groups.txt new file mode 100644 index 0000000..43b0daa --- /dev/null +++ b/configs/merge_groups.txt @@ -0,0 +1,94 @@ +# 병합 그룹 정의 — 대괄호 안이 병합 후 대표 라벨. +# 같은 그룹에 속한 라벨끼리 외곽선이 --gap px 이내로 인접하면 하나로 합친다. +# 어느 그룹에도 없는 라벨은 건드리지 않는다. + +# [building] +building +building rooftop +building facade +residential house +warehouse +factory structure +storage container +blue roof +red roof +green roof +gray roof +dark gray roof +black roof +white roof +silver metal roof +orange roof +brown roof +yellow roof +greenhouse roof +plastic greenhouse +vinyl greenhouse tunnel + +# [road] +asphalt road +concrete pavement +parking lot +pedestrian sidewalk +crosswalk + +# [road marking] +road lane marking +yellow center dividing line +white solid lane marking +white dashed lane marking +white directional arrow marking +yellow parking stall line +white parking stall line +blue handicap parking space +pink pedestrian safety marking +green bike lane marking + +# [vegetation] +tree canopy +dense forest +green hedge +grass lawn + +# [ground] +bare ground +dirt field + +# [stored material] +blue plastic drum +orange plastic barrel +yellow plastic container +white industrial tank +metallic storage tank +blue waterproof tarp +green waterproof tarp +black protective sheet +white canvas canopy + +# [fence] +gray metal fence +green wire fence +red roadside barrier + +# [vehicle] +white vehicle +black vehicle +silver vehicle +gray vehicle +red vehicle +blue vehicle +yellow vehicle +orange vehicle +green vehicle +white cargo truck +blue cargo truck +yellow school bus +commercial bus +tractor +farm tractor +excavator + +# 아래는 기본 비활성 — 나란히 놓인 콘이 한 덩어리로 융합된다. +# [cone] +# yellow safety cone +# orange traffic cone diff --git a/prompts/discovery_v1.txt b/prompts/discovery_v1.txt new file mode 100644 index 0000000..7082cf3 --- /dev/null +++ b/prompts/discovery_v1.txt @@ -0,0 +1,84 @@ +# SAM 3.1 Segmentation Prompts - Color & Object Combinations + +# [General Categories] +building +building rooftop +building facade +residential house +warehouse +storage container +factory structure +asphalt road +concrete pavement +parking lot +pedestrian sidewalk +road lane marking +crosswalk +bare ground +dirt field +grass lawn +tree canopy +dense forest +green hedge +cast shadow + +# [Roofs & Structures by Color] +blue roof +red roof +green roof +gray roof +dark gray roof +black roof +white roof +silver metal roof +orange roof +brown roof +yellow roof + +# [Vehicles & Transportation by Color] +white vehicle +black vehicle +silver vehicle +gray vehicle +red vehicle +blue vehicle +yellow vehicle +orange vehicle +green vehicle +white cargo truck +blue cargo truck +yellow school bus +commercial bus +tractor +farm tractor +excavator + +# [Industrial Materials, Storage & Objects by Color] +blue plastic drum +orange plastic barrel +yellow plastic container +white industrial tank +metallic storage tank +blue waterproof tarp +green waterproof tarp +black protective sheet +white canvas canopy +plastic greenhouse +vinyl greenhouse tunnel +greenhouse roof +gray metal fence +green wire fence +red roadside barrier +yellow safety cone +orange traffic cone + +# [Road Markings & Surface Features] +yellow center dividing line +white solid lane marking +white dashed lane marking +white directional arrow marking +yellow parking stall line +white parking stall line +blue handicap parking space +pink pedestrian safety marking +green bike lane marking diff --git a/prompts/wide_v1.txt b/prompts/wide_v1.txt new file mode 100644 index 0000000..505ca79 --- /dev/null +++ b/prompts/wide_v1.txt @@ -0,0 +1,12 @@ +# 타일보다 큰 대상 — 통짜(1×1)로 검출한다. +# 여기 적힌 라벨은 타일 패스에서 제외된다. + +building +asphalt road +concrete pavement +parking lot +bare ground +dirt field +grass lawn +tree canopy +dense forest diff --git a/tools/cut_tiles.py b/tools/cut_tiles.py new file mode 100644 index 0000000..143ab6b --- /dev/null +++ b/tools/cut_tiles.py @@ -0,0 +1,75 @@ +""" +이미지를 cols×rows 타일로 잘라 파일로 저장하고, R{행}C{열} 라벨을 찍은 격자 오버레이도 생성. +사용법: + python tools/cut_tiles.py --input <이미지> --cols 9 --rows 6 +출력: + output/tiles/<이미지명>/R1C1.jpg ... (원본 배율) + output/tiles/<이미지명>_grid.jpg (격자 오버레이, 최대 4096px) +""" +import argparse +from pathlib import Path + +import cv2 +import numpy as np + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--input", required=True) + ap.add_argument("--outdir", default="output/tiles") + ap.add_argument("--cols", type=int, default=9) + ap.add_argument("--rows", type=int, default=6) + args = ap.parse_args() + + img_path = Path(args.input) + buf = np.fromfile(str(img_path), dtype=np.uint8) + img = cv2.imdecode(buf, cv2.IMREAD_COLOR) + if img is None: + print(f"이미지 로드 실패: {img_path}") + return + + H, W = img.shape[:2] + base_w = W / args.cols + base_h = H / args.rows + print(f"이미지 {W}×{H} → {args.cols}×{args.rows} 타일 {base_w:.0f}×{base_h:.0f}") + + tile_dir = Path(args.outdir) / img_path.stem + tile_dir.mkdir(parents=True, exist_ok=True) + + vis = img.copy() + font_scale = base_h / 200.0 + thickness = max(2, int(font_scale * 2)) + + for r in range(args.rows): + for c in range(args.cols): + x0, x1 = int(c * base_w), int((c + 1) * base_w) + y0, y1 = int(r * base_h), int((r + 1) * base_h) + name = f"R{r+1}C{c+1}" + + cv2.imencode(".jpg", img[y0:y1, x0:x1], + [cv2.IMWRITE_JPEG_QUALITY, 95])[1].tofile( + str(tile_dir / f"{name}.jpg")) + + cv2.rectangle(vis, (x0, y0), (x1, y1), (0, 200, 255), 4) + (tw, th), _ = cv2.getTextSize(name, cv2.FONT_HERSHEY_SIMPLEX, + font_scale, thickness) + tx, ty = x0 + 12, y0 + th + 12 + cv2.rectangle(vis, (tx - 6, ty - th - 6), (tx + tw + 6, ty + 8), + (0, 0, 0), -1) + cv2.putText(vis, name, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX, + font_scale, (0, 200, 255), thickness, cv2.LINE_AA) + + 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))) + + grid_path = Path(args.outdir) / f"{img_path.stem}_grid.jpg" + cv2.imencode(".jpg", vis, [cv2.IMWRITE_JPEG_QUALITY, 92])[1].tofile(str(grid_path)) + + print(f"타일 {args.cols * args.rows}개 → {tile_dir}") + print(f"격자 오버레이 → {grid_path}") + + +if __name__ == "__main__": + main() diff --git a/tools/make_viewer.py b/tools/make_viewer.py new file mode 100644 index 0000000..bf6d6d5 --- /dev/null +++ b/tools/make_viewer.py @@ -0,0 +1,302 @@ +""" +세그멘테이션 결과 뷰어 생성 — 라벨별 on/off, 줌/팬. + +sam3_multi_prompt.py 가 만든 JSON과 원본 이미지를 받아 단독 HTML을 만든다. +라벨은 프롬프트 파일의 "# [그룹명]" 주석 기준으로 묶어서 나열한다. + +사용법: + python tools/make_viewer.py \ + --json "data/everyimage/output/0857_multi.json" \ + --image "data/everyimage/DJI_20250805162831_0857.JPG" \ + --prompts prompts/discovery_v1.txt +출력: + /<이름>_viewer.html + <이름>_view.jpg +""" +import argparse +import json +from collections import Counter, OrderedDict +from pathlib import Path + +import cv2 +import numpy as np + + +def parse_groups(path: Path) -> "OrderedDict[str, list]": + """프롬프트 파일의 '# [그룹명]' 주석으로 라벨을 묶는다.""" + groups, current = OrderedDict(), "기타" + for line in path.read_text(encoding="utf-8").splitlines(): + s = line.strip() + if s.startswith("#"): + body = s.lstrip("#").strip() + if body.startswith("[") and body.endswith("]"): + current = body[1:-1].strip() + groups.setdefault(current, []) + elif s: + groups.setdefault(current, []).append(s) + return groups + + +HTML = """ + +__TITLE__ + +
+

__TITLE__

+
segment __TOTAL__개 · 라벨 __NLAB__종
+
+ + + + + +
+
+
+
+ +""" + + +def color_for(label: str) -> str: + """라벨 문자열 해시 → 고정 색상.""" + h = 0 + for ch in label: + h = (h * 131 + ord(ch)) & 0xFFFFFFFF + hsv = np.uint8([[[(h % 360) // 2, # OpenCV 색상은 0~179 + 190 + (h >> 9) % 60, + 200 + (h >> 17) % 55]]]) + r, g, b = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)[0][0] + return "#%02x%02x%02x" % (int(r), int(g), int(b)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--json", required=True) + ap.add_argument("--image", required=True) + ap.add_argument("--prompts", default="prompts/discovery_v1.txt") + ap.add_argument("--output", default=None, help="기본: _viewer.html") + ap.add_argument("--max-size", type=int, default=3000, help="뷰어용 이미지 최대 변 길이") + args = ap.parse_args() + + jpath = Path(args.json) + data = json.loads(jpath.read_text(encoding="utf-8")) + segs = data.get("segments", []) + + img = cv2.imdecode(np.fromfile(args.image, dtype=np.uint8), cv2.IMREAD_COLOR) + if img is None: + print(f"이미지 로드 실패: {args.image}") + return + H, W = img.shape[:2] + k = min(1.0, args.max_size / max(H, W)) + if k < 1.0: + img = cv2.resize(img, (int(W * k), int(H * k)), interpolation=cv2.INTER_AREA) + + out_html = Path(args.output) if args.output else jpath.with_name(jpath.stem + "_viewer.html") + img_name = jpath.stem + "_view.jpg" + cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 88])[1].tofile( + str(out_html.with_name(img_name))) + + counts = Counter(s.get("label", "") for s in segs) + groups_src = parse_groups(Path(args.prompts)) + seen, groups = set(), [] + for name, labels in groups_src.items(): + rows = [{"label": lb, "color": color_for(lb), "count": counts.get(lb, 0)} + for lb in labels if counts.get(lb, 0) > 0] + seen.update(r["label"] for r in rows) + if rows: + groups.append({"name": name, "labels": rows}) + extra = [{"label": lb, "color": color_for(lb), "count": c} + for lb, c in counts.most_common() if lb not in seen and lb] + if extra: + groups.append({"name": "그룹 없음", "labels": extra}) + + payload = { + "groups": groups, + "colorOf": {lb: color_for(lb) for lb in counts if lb}, + "segs": [{"l": s.get("label", ""), + "p": [[round(p[0] * k, 1), round(p[1] * k, 1)] for p in s["points"]]} + for s in segs if s.get("points")], + } + + html = (HTML.replace("__TITLE__", jpath.stem) + .replace("__TOTAL__", str(len(segs))) + .replace("__NLAB__", str(len([c for c in counts if c]))) + .replace("__IMG__", img_name) + .replace("__DATA__", json.dumps(payload, ensure_ascii=False))) + out_html.write_text(html, encoding="utf-8") + + print(f"뷰어: {out_html}") + print(f"이미지: {out_html.with_name(img_name)} ({img.shape[1]}×{img.shape[0]})") + print(f"segment {len(segs)}개 · 라벨 {len([c for c in counts if c])}종 · 그룹 {len(groups)}개") + + +if __name__ == "__main__": + main() diff --git a/tools/merge_labels.py b/tools/merge_labels.py new file mode 100644 index 0000000..ecefd1c --- /dev/null +++ b/tools/merge_labels.py @@ -0,0 +1,214 @@ +""" +검출 결과 후처리 — 같은 병합 그룹에 속한 라벨끼리 외곽선이 gap px 이내로 +인접하면 하나로 합치고 그룹 대표 라벨을 붙인다. + +예) "building" + "building rooftop" 이 맞닿아 있으면 → "building" 하나로. + +사용법: + python tools/merge_labels.py \ + --json "data/everyimage/output/0006_multi.json" \ + --groups configs/merge_groups.txt \ + --gap 2 +출력: + <입력>_merged.json (뷰어에 그대로 넣을 수 있음) +""" +import argparse +import json +import sys +from collections import Counter, OrderedDict +from pathlib import Path + +import cv2 +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from sam3_everything_explore import _bbox, _polys_touch # noqa: E402 + + +def load_merge_groups(path: Path) -> "OrderedDict[str, str]": + """라벨 → 대표 라벨 매핑. '# [대표]' 헤더 아래 라벨들이 그 그룹.""" + mapping, current = OrderedDict(), None + for line in path.read_text(encoding="utf-8").splitlines(): + s = line.strip() + if not s: + continue + if s.startswith("#"): + body = s.lstrip("#").strip() + current = body[1:-1].strip() if body.startswith("[") and body.endswith("]") else None + continue + if current: + mapping[s] = current + return mapping + + +def load_claim_rules(path: Path): + """'<그룹> : <비율> : <대상 라벨들>' → [(그룹, 비율, {대상라벨})]""" + rules = [] + for line in path.read_text(encoding="utf-8").splitlines(): + s = line.strip() + if not s or s.startswith("#"): + continue + parts = [p.strip() for p in s.split(":")] + if len(parts) != 3: + raise SystemExit(f"claim 규칙 형식 오류: {line}") + rules.append((parts[0], float(parts[1]), + {t.strip() for t in parts[2].split(",") if t.strip()})) + return rules + + +def _overlap_ratio(pa, pb): + """pa 면적 대비 pa∩pb 비율.""" + xs = [p[0] for p in pa] + [p[0] for p in pb] + ys = [p[1] for p in pa] + [p[1] for p in pb] + x0, y0 = int(min(xs)) - 1, int(min(ys)) - 1 + x1, y1 = int(max(xs)) + 1, int(max(ys)) + 1 + ca = np.zeros((y1 - y0, x1 - x0), np.uint8) + cb = np.zeros_like(ca) + cv2.fillPoly(ca, [np.array(pa, np.int32) - (x0, y0)], 255) + cv2.fillPoly(cb, [np.array(pb, np.int32) - (x0, y0)], 255) + area = int(np.count_nonzero(ca)) + return 0.0 if area == 0 else np.count_nonzero(cv2.bitwise_and(ca, cb)) / area + + +def claim_labels(shapes, mapping, rules): + """겹침 기준으로 라벨 소유권을 재배정한다. 반환: 바뀐 개수.""" + changed = 0 + boxes = [_bbox(s["points"]) for s in shapes] + for group, ratio, targets in rules: + owners = [i for i, s in enumerate(shapes) + if mapping.get(s.get("label", "")) == group] + if not owners: + print(f" [claim] 그룹 '{group}' 폴리곤 없음 — 규칙 무시") + continue + for i, s in enumerate(shapes): + if s.get("label", "") not in targets: + continue + ax0, ay0, ax1, ay1 = boxes[i] + best, best_r = None, 0.0 + for j in owners: + bx0, by0, bx1, by1 = boxes[j] + if ax1 < bx0 or bx1 < ax0 or ay1 < by0 or by1 < ay0: + continue + r = _overlap_ratio(s["points"], shapes[j]["points"]) + if r > best_r: + best, best_r = shapes[j], r + if best is not None and best_r >= ratio: + s["claimed_from"] = s["label"] + s["label"] = best["label"] + changed += 1 + return changed + + +def merge_by_group(shapes, mapping, gap=2, epsilon=1.5): + """대표 라벨이 같은 것끼리 외곽선 인접 시 병합. 그룹 밖 라벨은 그대로 통과.""" + grouped, passthrough = {}, [] + for s in shapes: + rep = mapping.get(s.get("label", "")) + if rep is None: + passthrough.append(s) + else: + grouped.setdefault(rep, []).append(s) + + merged = list(passthrough) + for rep, items in grouped.items(): + parent = list(range(len(items))) + + def find(i): + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + boxes = [_bbox(s["points"]) for s in items] + for i in range(len(items)): + for j in range(i + 1, len(items)): + if find(i) == find(j): + continue + ax0, ay0, ax1, ay1 = boxes[i] + bx0, by0, bx1, by1 = boxes[j] + if ax1 + gap < bx0 or bx1 + gap < ax0 or ay1 + gap < by0 or by1 + gap < ay0: + continue # 조기 탈락 (판정은 아래 픽셀 단위) + if _polys_touch(items[i]["points"], items[j]["points"], gap): + parent[find(j)] = find(i) + + clusters = {} + for i in range(len(items)): + clusters.setdefault(find(i), []).append(i) + + for members in clusters.values(): + src = [items[i] for i in members] + best = max(src, key=lambda s: float(s.get("score", 0))) + if len(src) == 1: + merged.append({**src[0], "label": rep, + "merged_from": [src[0].get("label", "")]}) + continue + pts = [p for s in src for p in s["points"]] + x0 = int(min(p[0] for p in pts)) - gap - 1 + y0 = int(min(p[1] for p in pts)) - gap - 1 + x1 = int(max(p[0] for p in pts)) + gap + 1 + y1 = int(max(p[1] for p in pts)) + gap + 1 + canvas = np.zeros((y1 - y0, x1 - x0), np.uint8) + for s in src: + cv2.fillPoly(canvas, [np.array(s["points"], np.int32) - (x0, y0)], 255) + if gap > 0: # 틈 메우기 + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (gap * 2 + 1, gap * 2 + 1)) + canvas = cv2.morphologyEx(canvas, cv2.MORPH_CLOSE, k) + contours, _ = cv2.findContours(canvas, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + names = sorted({s.get("label", "") for s in src}) + for cnt in contours: + approx = cv2.approxPolyDP(cnt, epsilon, True) + if len(approx) < 3: + continue + merged.append({ + "label": rep, + "score": float(best.get("score", 0)), + "shape_type": "polygon", + "merged_from": names, + "points": [[float(p[0][0] + x0), float(p[0][1] + y0)] for p in approx], + }) + return merged + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--json", required=True) + ap.add_argument("--groups", default="configs/merge_groups.txt") + ap.add_argument("--claims", default="configs/claim_rules.txt", + help="소유권 재배정 규칙 (없으면 빈 문자열로 끄기)") + ap.add_argument("--gap", type=int, default=2, help="외곽선 인접 판정 px (기본 2)") + ap.add_argument("--output", default=None, help="기본: <입력>_merged.json") + args = ap.parse_args() + + jpath = Path(args.json) + data = json.loads(jpath.read_text(encoding="utf-8")) + shapes = data.get("segments", []) + mapping = load_merge_groups(Path(args.groups)) + + reps = sorted(set(mapping.values())) + print(f"입력 {len(shapes)}개 · 병합 그룹 {len(reps)}개: {', '.join(reps)}") + + if args.claims: + rules = load_claim_rules(Path(args.claims)) + n = claim_labels(shapes, mapping, rules) + print(f"소유권 재배정: {n}개 라벨 변경") + + out = merge_by_group(shapes, mapping, gap=args.gap) + counts = Counter(s.get("label", "") for s in out) + print(f"병합(gap={args.gap}px) {len(shapes)} → {len(out)}개\n") + for lb, c in counts.most_common(): + print(f" {lb:34s} {c:4d}") + + out_path = Path(args.output) if args.output else jpath.with_name(jpath.stem + "_merged.json") + out_path.write_text(json.dumps({ + "total_segments": len(out), + "label_counts": dict(counts), + "segments": [{"label": s.get("label", ""), "score": s.get("score", 0), + "merged_from": s.get("merged_from", []), + "bbox": list(_bbox(s["points"])), "points": s["points"]} + for s in out], + }, ensure_ascii=False, indent=2), encoding="utf-8") + print(f"\n저장: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/tools/sam3_everything_explore.py b/tools/sam3_everything_explore.py index 929b7fc..8bdddc4 100644 --- a/tools/sam3_everything_explore.py +++ b/tools/sam3_everything_explore.py @@ -39,26 +39,34 @@ def encode_image(image_bgr: np.ndarray, max_size: int = 1280) -> tuple: return base64.b64encode(buf).decode("utf-8"), scale -# 탐색용 넓은 프롬프트 — 철도 현장에서 흔히 보이는 모든 요소 포함 +# 탐색용 프롬프트 — 색상·객체 조합 67개 DISCOVERY_PROMPT = ( - "railroad track, railway rail, " - "catenary pole, overhead line pole, electric pole, " - "overhead wire, catenary wire, power line cable, " - "railway sleeper, concrete tie, " - "guardrail, highway barrier, road fence, " - "bridge, viaduct, overpass, " - "vegetation, tree, bush, grass, " - "building, structure, roof, wall, " - "vehicle, car, truck, " - "road, asphalt, pavement, " - "slope, embankment, retaining wall, " - "noise barrier, sound wall, " - "signal, sign board, " - "small dark object on ballast, small dark object on railway, " - "small square metal box on ground, control box on ballast, " - "gray square lid on gravel, flat metal cover on ground, " - "small bright object on ballast, small white box on ballast, " - "small gray box on ground, bright square object on gravel" + # 일반 카테고리 + "building, building rooftop, building facade, residential house, " + "warehouse, storage container, factory structure, " + "asphalt road, concrete pavement, parking lot, pedestrian sidewalk, " + "road lane marking, crosswalk, " + "bare ground, dirt field, grass lawn, " + "tree canopy, dense forest, green hedge, cast shadow, " + # 지붕·구조물 (색상별) + "blue roof, red roof, green roof, gray roof, dark gray roof, black roof, " + "white roof, silver metal roof, orange roof, brown roof, yellow roof, " + # 차량 (색상별) + "white vehicle, black vehicle, silver vehicle, gray vehicle, red vehicle, " + "blue vehicle, yellow vehicle, orange vehicle, green vehicle, " + "white cargo truck, blue cargo truck, yellow school bus, commercial bus, " + # 산업 자재·적치물 (색상별) + "blue plastic drum, orange plastic barrel, yellow plastic container, " + "white industrial tank, metallic storage tank, " + "blue waterproof tarp, green waterproof tarp, black protective sheet, " + "white canvas canopy, gray metal fence, green wire fence, " + "red roadside barrier, yellow safety cone, orange traffic cone, " + # 노면 표시 + "yellow center dividing line, white solid lane marking, " + "white dashed lane marking, white directional arrow marking, " + "yellow parking stall line, white parking stall line, " + "blue handicap parking space, pink pedestrian safety marking, " + "green bike lane marking" ) @@ -178,6 +186,84 @@ def detect_everything_tiled(image_bgr, cols, rows, overlap, conf, workers, promp return all_shapes +# ── 인접 폴리곤 병합 ────────────────────────────────────────────────────────── +def _polys_touch(pa, pb, gap): + """두 폴리곤이 gap px 이내로 닿는지 (겹침 포함) 국소 캔버스에서 판정.""" + xs = [p[0] for p in pa] + [p[0] for p in pb] + ys = [p[1] for p in pa] + [p[1] for p in pb] + x0, y0 = int(min(xs)) - gap - 1, int(min(ys)) - gap - 1 + x1, y1 = int(max(xs)) + gap + 1, int(max(ys)) + gap + 1 + if (x1 - x0) * (y1 - y0) > 40_000_000: + return False + ca = np.zeros((y1 - y0, x1 - x0), np.uint8) + cb = np.zeros_like(ca) + cv2.fillPoly(ca, [np.array(pa, np.int32) - (x0, y0)], 255) + cv2.fillPoly(cb, [np.array(pb, np.int32) - (x0, y0)], 255) + if gap > 0: + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (gap * 2 + 1, gap * 2 + 1)) + ca = cv2.dilate(ca, k) + return bool(np.any(cv2.bitwise_and(ca, cb))) + + +def merge_adjacent(shapes, gap=8, epsilon=2.0): + """같은 라벨끼리 gap px 이내로 인접·겹치는 폴리곤을 하나로 병합.""" + if not shapes: + return [] + + parent = list(range(len(shapes))) + + def find(i): + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + boxes = [_bbox(s["points"]) for s in shapes] + for i in range(len(shapes)): + for j in range(i + 1, len(shapes)): + if shapes[i].get("label") != shapes[j].get("label"): + continue + if find(i) == find(j): + continue + ax0, ay0, ax1, ay1 = boxes[i] + bx0, by0, bx1, by1 = boxes[j] + if ax1 + gap < bx0 or bx1 + gap < ax0 or ay1 + gap < by0 or by1 + gap < ay0: + continue # bbox조차 안 닿음 + if _polys_touch(shapes[i]["points"], shapes[j]["points"], gap): + parent[find(j)] = find(i) + + groups = {} + for i in range(len(shapes)): + groups.setdefault(find(i), []).append(i) + + merged = [] + for members in groups.values(): + if len(members) == 1: + merged.append(shapes[members[0]]) + continue + pts_all = [p for i in members for p in shapes[i]["points"]] + x0 = int(min(p[0] for p in pts_all)) - gap - 1 + y0 = int(min(p[1] for p in pts_all)) - gap - 1 + x1 = int(max(p[0] for p in pts_all)) + gap + 1 + y1 = int(max(p[1] for p in pts_all)) + gap + 1 + canvas = np.zeros((y1 - y0, x1 - x0), np.uint8) + for i in members: + cv2.fillPoly(canvas, [np.array(shapes[i]["points"], np.int32) - (x0, y0)], 255) + if gap > 0: + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (gap * 2 + 1, gap * 2 + 1)) + canvas = cv2.morphologyEx(canvas, cv2.MORPH_CLOSE, k) + contours, _ = cv2.findContours(canvas, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + base = shapes[max(members, key=lambda i: float(shapes[i].get("score", 0)))] + for cnt in contours: + approx = cv2.approxPolyDP(cnt, epsilon, True) + if len(approx) < 3: + continue + merged.append({**base, + "points": [[float(p[0][0] + x0), float(p[0][1] + y0)] + for p in approx]}) + return merged + + # ── 시각화 ──────────────────────────────────────────────────────────────────── def draw_everything(image_bgr, shapes, cols, rows): vis = image_bgr.copy() @@ -190,6 +276,10 @@ def draw_everything(image_bgr, shapes, cols, rows): bx1, by1 = int((c + 1) * W / cols), int((r + 1) * H / rows) cv2.rectangle(vis, (bx0, by0), (bx1, by1), (60, 60, 60), 1) + # 라벨 글자 크기는 이미지 크기 비례 (출력이 4096으로 축소되는 것 감안) + font_scale = max(0.8, min(W, H) / 2200) + font_thick = max(2, int(font_scale * 1.5)) + rng = np.random.default_rng(42) for s in shapes: pts = np.array(s["points"], dtype=np.int32) @@ -197,15 +287,26 @@ def draw_everything(image_bgr, shapes, cols, rows): overlay = vis.copy() cv2.fillPoly(overlay, [pts], color) cv2.addWeighted(overlay, 0.30, vis, 0.70, 0, vis) - cv2.polylines(vis, [pts], True, color, 1) + + # 외곽선: 선을 그린 뒤 마스크 내부만 남겨 안쪽 1px로 만듦 + region = np.zeros((H, W), dtype=np.uint8) + cv2.fillPoly(region, [pts], 255) + edge = np.zeros((H, W), dtype=np.uint8) + cv2.polylines(edge, [pts], True, 255, 2, cv2.LINE_8) + vis[cv2.bitwise_and(edge, region) > 0] = [int(c * 0.35) for c in color] # 라벨 표시 (있을 경우) label = s.get("label", "") if label: cx = int(np.mean([p[0] for p in s["points"]])) cy = int(np.mean([p[1] for p in s["points"]])) - cv2.putText(vis, label, (cx, cy), - cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1, cv2.LINE_AA) + (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, + font_scale, font_thick) + cv2.rectangle(vis, (cx - tw // 2 - 4, cy - th - 4), + (cx + tw // 2 + 4, cy + 6), (0, 0, 0), -1) + cv2.putText(vis, label, (cx - tw // 2, cy), + cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, + font_thick, cv2.LINE_AA) cv2.putText(vis, f"total segments: {len(shapes)}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2) @@ -246,6 +347,8 @@ def main(): ap.add_argument("--workers", type=int, default=4, help="병렬 스레드 수 (기본 4)") ap.add_argument("--nms", type=float, default=0.40, help="NMS IoU 임계값 (기본 0.40)") ap.add_argument("--prompt-extra", default="", help="DISCOVERY_PROMPT 뒤에 추가할 어휘 (콤마 구분)") + ap.add_argument("--merge", action="store_true", help="같은 라벨끼리 인접 폴리곤 병합") + ap.add_argument("--merge-gap", type=int, default=8, help="병합 판정 간격 px (기본 8)") ap.add_argument("--zone", type=int, nargs=4, metavar=("X1","Y1","X2","Y2"), default=None, help="처리 zone 제한 (이 범위와 겹치는 타일만 처리)") args = ap.parse_args() @@ -280,7 +383,13 @@ def main(): ) print(f"검출 {len(shapes)}개 → NMS(iou={args.nms})...") shapes = nms_shapes(shapes, iou_thresh=args.nms) - print(f"NMS 후 {len(shapes)}개 ({time.time()-t0:.0f}초)\n") + print(f"NMS 후 {len(shapes)}개 ({time.time()-t0:.0f}초)") + + if args.merge: + before = len(shapes) + shapes = merge_adjacent(shapes, gap=args.merge_gap) + print(f"병합(gap={args.merge_gap}px) {before} → {len(shapes)}개") + print() analyze_labels(shapes) diff --git a/tools/sam3_multi_prompt.py b/tools/sam3_multi_prompt.py new file mode 100644 index 0000000..90a0972 --- /dev/null +++ b/tools/sam3_multi_prompt.py @@ -0,0 +1,298 @@ +""" +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() diff --git a/tools/sam3_segment_everything.py b/tools/sam3_segment_everything.py index 12f6dfb..46d2930 100644 --- a/tools/sam3_segment_everything.py +++ b/tools/sam3_segment_everything.py @@ -8,16 +8,17 @@ import numpy as np import torch # Add server to path so we can import sam3 locally -server_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "X-AnyLabeling-Server")) -models_path = os.path.join(server_path, "app", "models") +server_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "sam31server")) if server_path not in sys.path: sys.path.insert(0, server_path) -if models_path not in sys.path: - sys.path.insert(0, models_path) - + +SAM3_CHECKPOINT = ( + "C:/Users/nbright/.cache/huggingface/hub/models--facebook--sam3.1/" + "snapshots/daa63191845a41281374e725f4c9e51c7a824460/sam3.1_multiplex.pt" +) + from sam3.model_builder import build_sam3_image_model from sam3.model.sam3_image_processor import Sam3Processor -from app.models.segment_anything_3 import SegmentAnything3 def build_point_grid(n_per_side: int) -> np.ndarray: """Generates a 2D grid of points evenly spaced in [0, 1] x [0, 1].""" @@ -142,6 +143,10 @@ def main(): parser.add_argument("--input", required=True, help="Input image path") parser.add_argument("--output", required=True, help="Output vis image path") parser.add_argument("--points", type=int, default=32, help="Points per side") + parser.add_argument("--conf", type=float, default=0.7, help="Confidence threshold") + parser.add_argument("--nms", type=float, default=0.7, help="NMS IoU threshold") + parser.add_argument("--split", action="store_true", + help="겹치지 않는 마스크끼리 묶어 여러 장으로 저장") args = parser.parse_args() buf = np.fromfile(args.input, dtype=np.uint8) @@ -159,27 +164,65 @@ def main(): else: image_proc = image.copy() - model_path = os.path.join(server_path, "sam3.pt") + model_path = SAM3_CHECKPOINT - results = segment_everything(image_proc, model_path, points_per_side=args.points, conf_thresh=0.7, nms_thresh=0.7) + results = segment_everything(image_proc, model_path, points_per_side=args.points, + conf_thresh=args.conf, nms_thresh=args.nms) - vis = image_proc.copy() np.random.seed(42) + colors = [np.random.randint(0, 255, (3,)).tolist() for _ in results] + H, W = image_proc.shape[:2] + regions = [] for res in results: - poly = res["polygon"] - pts = np.array(poly, dtype=np.int32) - color = np.random.randint(0, 255, (3,)).tolist() - - overlay = vis.copy() - cv2.fillPoly(overlay, [pts], color) - cv2.addWeighted(overlay, 0.4, vis, 0.6, 0, vis) - cv2.polylines(vis, [pts], True, color, 1) - - # Fix unicode paths in output - is_success, im_buf_arr = cv2.imencode(".jpg", vis) - if is_success: - im_buf_arr.tofile(args.output) - print(f"Saved visualization to {args.output}") + region = np.zeros((H, W), dtype=np.uint8) + cv2.fillPoly(region, [np.array(res["polygon"], dtype=np.int32)], 255) + regions.append(region) + + if args.split: + # 겹치지 않는 것끼리 묶어 레이어 분리 (점수 높은 순 greedy first-fit) + layers = [] # [(누적마스크, [인덱스...])] + for i, region in enumerate(regions): + for canvas, members in layers: + if not np.any(cv2.bitwise_and(canvas, region)): + cv2.bitwise_or(canvas, region, canvas) + members.append(i) + break + else: + layers.append((region.copy(), [i])) + print(f"Split into {len(layers)} non-overlapping layers.") + groups = [members for _, members in layers] + else: + groups = [list(range(len(results)))] + + out_path = Path(args.output) + for n, members in enumerate(groups, 1): + vis = image_proc.copy() + for i in members: + pts = np.array(results[i]["polygon"], dtype=np.int32) + color = colors[i] + + overlay = vis.copy() + cv2.fillPoly(overlay, [pts], color) + cv2.addWeighted(overlay, 0.4, vis, 0.6, 0, vis) + + # 테두리: 선을 그린 뒤 마스크 내부만 남겨 안쪽 1px로 만듦 + edge = np.zeros((H, W), dtype=np.uint8) + cv2.polylines(edge, [pts], True, 255, 2, cv2.LINE_8) + vis[cv2.bitwise_and(edge, regions[i]) > 0] = [int(c * 0.35) for c in color] + + path = (out_path if len(groups) == 1 + else out_path.with_name(f"{out_path.stem}_L{n}{out_path.suffix}")) + is_success, im_buf_arr = cv2.imencode(out_path.suffix, vis) + if is_success: + im_buf_arr.tofile(str(path)) + print(f"Saved {len(members)} masks to {path}") + + try: # 완료 알림음 + import winsound + winsound.Beep(880, 150) + winsound.Beep(1175, 250) + except Exception: + print("\a", end="") if __name__ == "__main__": main()