yolo最新选择
This commit is contained in:
yrx
@@ -0,0 +1,184 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user