91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
"""
|
||
6×8 (또는 임의 그리드) 타일 라벨을 원본 이미지 좌표로 병합해 시각화.
|
||
|
||
사용:
|
||
python tools/merge_tiles_vis.py \
|
||
--orig data/역사이미지/slope/DJI_20260306113838_0004.JPG \
|
||
--labels output/autolabel/tile6x8/labels \
|
||
--output output/autolabel/tile6x8/merged_vis.jpg \
|
||
--cols 6 --rows 8
|
||
"""
|
||
import argparse
|
||
import cv2
|
||
import numpy as np
|
||
from pathlib import Path
|
||
|
||
CLASS_NAMES = ["catenary_pole", "bracket"]
|
||
CLASS_COLORS = [(0, 200, 255), (255, 130, 0)]
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--orig", required=True)
|
||
parser.add_argument("--labels", required=True)
|
||
parser.add_argument("--output", required=True)
|
||
parser.add_argument("--cols", type=int, default=6)
|
||
parser.add_argument("--rows", type=int, default=8)
|
||
parser.add_argument("--alpha", type=float, default=0.3)
|
||
args = parser.parse_args()
|
||
|
||
buf = np.fromfile(args.orig, dtype=np.uint8)
|
||
img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
|
||
H, W = img.shape[:2]
|
||
tw, th = W // args.cols, H // args.rows
|
||
|
||
label_dir = Path(args.labels)
|
||
counts = [0] * len(CLASS_NAMES)
|
||
|
||
for r in range(args.rows):
|
||
for c in range(args.cols):
|
||
label_file = label_dir / f"tile_r{r+1:02d}_c{c+1:02d}.txt"
|
||
if not label_file.exists():
|
||
continue
|
||
x0 = c * tw
|
||
y0 = r * th
|
||
tile_w = tw if c < args.cols - 1 else W - x0
|
||
tile_h = th if r < args.rows - 1 else H - y0
|
||
|
||
text = label_file.read_text(encoding="utf-8").strip()
|
||
if not text:
|
||
continue
|
||
for line in text.splitlines():
|
||
parts = line.split()
|
||
if not parts:
|
||
continue
|
||
cls_id = int(parts[0])
|
||
coords = list(map(float, parts[1:]))
|
||
pts = np.array(
|
||
[[coords[i] * tile_w + x0, coords[i + 1] * tile_h + y0]
|
||
for i in range(0, len(coords), 2)],
|
||
dtype=np.int32,
|
||
)
|
||
color = CLASS_COLORS[cls_id % len(CLASS_COLORS)]
|
||
overlay = img.copy()
|
||
cv2.fillPoly(overlay, [pts], color)
|
||
cv2.addWeighted(overlay, args.alpha, img, 1 - args.alpha, 0, img)
|
||
cv2.polylines(img, [pts], True, color, 3)
|
||
if cls_id < len(counts):
|
||
counts[cls_id] += 1
|
||
|
||
# 범례
|
||
for i, name in enumerate(CLASS_NAMES):
|
||
y = 30 + i * 50
|
||
cv2.rectangle(img, (15, y - 20), (55, y + 10), CLASS_COLORS[i], -1)
|
||
cv2.putText(img, f"{name}: {counts[i]}",
|
||
(65, y), cv2.FONT_HERSHEY_SIMPLEX, 1.5, CLASS_COLORS[i], 3)
|
||
|
||
total = sum(counts)
|
||
print(f"총 {total}개 " + ", ".join(f"{CLASS_NAMES[i]}={counts[i]}" for i in range(len(counts))))
|
||
|
||
# 4096px 이하로 미리보기 저장
|
||
scale = min(1.0, 4096 / max(H, W))
|
||
vis = cv2.resize(img, (int(W * scale), int(H * scale)))
|
||
out_path = Path(args.output)
|
||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
cv2.imencode(out_path.suffix, vis)[1].tofile(str(out_path))
|
||
print(f"저장: {out_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|