#!/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()