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()
|