185 lines
8.0 KiB
Python
185 lines
8.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""Offline baseline for traditional target-paper detection.
|
|
|
|
Dataset format: sibling .txt files use YOLO boxes and classes.txt maps ids
|
|
(the supplied dataset uses 0=40, 1=20, 2=10). This intentionally simple
|
|
baseline uses grayscale segmentation and contour geometry; it is useful as a
|
|
reference before adding more specialized black-triangle grouping.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import glob
|
|
import itertools
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def detect_white_papers(image: np.ndarray) -> list[tuple[int, int, int, int]]:
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
h, w = gray.shape[:2]
|
|
mask = cv2.inRange(gray, 120, 255)
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((9, 9), np.uint8))
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
out = []
|
|
for contour in contours:
|
|
x, y, bw, bh = cv2.boundingRect(contour)
|
|
area = float(bw * bh)
|
|
if area < 0.05 * w * h or min(bw, bh) < 80:
|
|
continue
|
|
fill = cv2.contourArea(contour) / max(area, 1.0)
|
|
aspect = bw / max(float(bh), 1.0)
|
|
if fill >= 0.45 and 0.4 <= aspect <= 2.5:
|
|
out.append((x, y, x + bw, y + bh))
|
|
return out
|
|
|
|
|
|
def detect_black_triangle_papers(image: np.ndarray):
|
|
"""Infer paper boxes from the four small black corner marks."""
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
mask = cv2.inRange(gray, 0, 100)
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8))
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
points = []
|
|
for contour in contours:
|
|
x, y, bw, bh = cv2.boundingRect(contour)
|
|
area = cv2.contourArea(contour)
|
|
vertices = cv2.approxPolyDP(contour, 0.08 * cv2.arcLength(contour, True), True)
|
|
if 60 <= area <= 400 and 8 <= bw <= 24 and 8 <= bh <= 24:
|
|
if 3 <= len(vertices) <= 5 and 0.5 <= bw / max(bh, 1) <= 2.0:
|
|
points.append((x + bw / 2.0, y + bh / 2.0))
|
|
candidates = []
|
|
for group in itertools.combinations(points, 4):
|
|
xs = sorted(p[0] for p in group)
|
|
ys = sorted(p[1] for p in group)
|
|
span_x, span_y = xs[-1] - xs[0], ys[-1] - ys[0]
|
|
if span_x < 50 or span_y < 50 or not 0.45 < span_x / span_y < 1.5:
|
|
continue
|
|
corners = ((xs[0], ys[0]), (xs[-1], ys[0]),
|
|
(xs[0], ys[-1]), (xs[-1], ys[-1]))
|
|
error = max(min(np.hypot(p[0] - c[0], p[1] - c[1]) for c in corners)
|
|
for p in group) / max(span_x, span_y)
|
|
if error > 0.22:
|
|
continue
|
|
ex, ey = 0.12 * span_x, 0.12 * span_y
|
|
candidates.append((xs[0] - ex, ys[0] - ey,
|
|
xs[-1] + ex, ys[-1] + ey, error))
|
|
# A colored target ring supplies an independent center check. Hough is
|
|
# deliberately low-cost here because it runs only on the already small
|
|
# candidate list's source frame.
|
|
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
|
color = cv2.inRange(hsv, (0, 70, 45), (179, 255, 255))
|
|
color = cv2.morphologyEx(color, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
|
|
ring_centers = []
|
|
for contour in cv2.findContours(color, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]:
|
|
area = cv2.contourArea(contour)
|
|
if area < 150:
|
|
continue
|
|
moments = cv2.moments(contour)
|
|
if moments["m00"]:
|
|
ring_centers.append((moments["m10"] / moments["m00"], moments["m01"] / moments["m00"]))
|
|
checked = []
|
|
for box in candidates:
|
|
if not ring_centers:
|
|
checked.append(box)
|
|
continue
|
|
x0, y0, x1, y1, err = box
|
|
inside = any(x0 - .15 * (x1 - x0) <= cx <= x1 + .15 * (x1 - x0)
|
|
and y0 - .15 * (y1 - y0) <= cy <= y1 + .15 * (y1 - y0)
|
|
for cx, cy in ring_centers)
|
|
if inside:
|
|
checked.append(box)
|
|
return sorted(checked, key=lambda x: x[-1])
|
|
|
|
|
|
def iou(a, b):
|
|
x0, y0 = max(a[0], b[0]), max(a[1], b[1])
|
|
x1, y1 = min(a[2], b[2]), min(a[3], b[3])
|
|
inter = max(0, x1 - x0) * max(0, y1 - y0)
|
|
aa = max(0, a[2] - a[0]) * max(0, a[3] - a[1])
|
|
bb = max(0, b[2] - b[0]) * max(0, b[3] - b[1])
|
|
return inter / max(aa + bb - inter, 1)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("dataset", help="directory containing jpg and YOLO txt files")
|
|
ap.add_argument("--iou", type=float, default=0.5)
|
|
ap.add_argument("--out", default="traditional_eval_results.csv",
|
|
help="CSV output path; relative paths are next to the dataset")
|
|
ap.add_argument("--vis-dir", default="traditional_eval_images",
|
|
help="directory for annotated result images; empty disables")
|
|
args = ap.parse_args()
|
|
stats = {0: [0, 0], 1: [0, 0]}
|
|
rows = []
|
|
# OpenCV on some Windows builds cannot decode non-ASCII filenames. Work
|
|
# relative to the dataset directory so the supplied Chinese path is safe.
|
|
dataset = os.path.abspath(args.dataset)
|
|
os.chdir(dataset)
|
|
# cwd is now the dataset, so a relative output avoids Windows console
|
|
# encoding issues with the Chinese parent path.
|
|
vis_dir = args.vis_dir if args.vis_dir else ""
|
|
if vis_dir:
|
|
os.makedirs(vis_dir, exist_ok=True)
|
|
files = glob.glob(os.path.join("**", "*.jpg"), recursive=True)
|
|
for image_path in files:
|
|
label_path = os.path.splitext(image_path)[0] + ".txt"
|
|
if not os.path.isfile(label_path):
|
|
continue
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
continue
|
|
h, w = image.shape[:2]
|
|
predictions = detect_black_triangle_papers(image)
|
|
vis = image.copy()
|
|
for p in predictions:
|
|
cv2.rectangle(vis, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), (0, 255, 255), 2)
|
|
for line in open(label_path, encoding="utf-8", errors="ignore"):
|
|
z = line.split()
|
|
if len(z) < 5 or int(float(z[0])) not in stats:
|
|
continue
|
|
cls, cx, cy, bw, bh = int(float(z[0])), *map(float, z[1:5])
|
|
truth = (int((cx - bw / 2) * w), int((cy - bh / 2) * h),
|
|
int((cx + bw / 2) * w), int((cy + bh / 2) * h))
|
|
best = max((iou(truth, p) for p in predictions), default=0.0)
|
|
best_box = max(predictions, key=lambda p: iou(truth, p), default=())
|
|
stats[cls][0] += 1
|
|
stats[cls][1] += int(best >= args.iou)
|
|
rows.append({
|
|
"image": image_path,
|
|
"class_id": cls,
|
|
"truth_xyxy": ",".join(map(str, truth[:4])),
|
|
"pred_xyxy": ",".join(map(str, best_box[:4])) if best_box else "",
|
|
"iou": f"{best:.4f}",
|
|
"pass": int(best >= args.iou),
|
|
})
|
|
color = (0, 255, 0) if best >= args.iou else (0, 0, 255)
|
|
cv2.rectangle(vis, truth[:2], truth[2:4], color, 2)
|
|
cv2.putText(vis, f"GT {cls} IoU {best:.2f}",
|
|
(truth[0], max(16, truth[1] - 4)),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1, cv2.LINE_AA)
|
|
if vis_dir:
|
|
name = os.path.splitext(os.path.basename(image_path))[0] + "_result.jpg"
|
|
cv2.imwrite(os.path.join(vis_dir, name), vis)
|
|
total = sum(v[0] for v in stats.values())
|
|
good = sum(v[1] for v in stats.values())
|
|
print(f"paper objects: {good}/{total} = {good / max(total, 1):.2%} (IoU >= {args.iou})")
|
|
for cls, (n, ok) in stats.items():
|
|
print(f"class {cls}: {ok}/{n} = {ok / max(n, 1):.2%}")
|
|
out_path = args.out if os.path.isabs(args.out) else os.path.join(dataset, args.out)
|
|
with open(out_path, "w", newline="", encoding="utf-8-sig") as fp:
|
|
writer = csv.DictWriter(fp, fieldnames=("image", "class_id", "truth_xyxy",
|
|
"pred_xyxy", "iou", "pass"))
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
print(f"details csv: {out_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|