69 lines
3.3 KiB
Python
69 lines
3.3 KiB
Python
import csv
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
weights = Path(r"D:\work\yolo\runs\detect\smoke_fire_yolo11s_v1-4\weights\best.pt")
|
|
candidates_path = Path(r"D:\work\yolo\runs\audit\suspected_missing_labels.csv")
|
|
output_dir = Path(r"D:\work\yolo\runs\audit\missing_label_review_ge_0.50")
|
|
threshold = 0.50
|
|
class_names = ["smoke", "fire"]
|
|
|
|
with candidates_path.open(encoding="utf-8", newline="") as csv_file:
|
|
candidates = list(csv.DictReader(csv_file))
|
|
image_paths = []
|
|
for row in candidates:
|
|
if float(row["confidence"]) < threshold:
|
|
continue
|
|
image_path = Path(row["image"])
|
|
if image_path not in image_paths:
|
|
image_paths.append(image_path)
|
|
|
|
model = YOLO(str(weights))
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
manifest = []
|
|
for index, image_path in enumerate(image_paths, start=1):
|
|
image = cv2.imread(str(image_path))
|
|
if image is None:
|
|
continue
|
|
height, width = image.shape[:2]
|
|
label_path = image_path.parents[1] / "labels" / f"{image_path.stem}.txt"
|
|
labeled = image.copy()
|
|
label_lines = [line for line in label_path.read_text(encoding="utf-8").splitlines() if line.strip()] if label_path.is_file() else []
|
|
for line in label_lines:
|
|
values = line.split()
|
|
class_id = int(float(values[0]))
|
|
x_center, y_center, box_width, box_height = map(float, values[1:])
|
|
x1 = int((x_center - box_width / 2) * width)
|
|
y1 = int((y_center - box_height / 2) * height)
|
|
x2 = int((x_center + box_width / 2) * width)
|
|
y2 = int((y_center + box_height / 2) * height)
|
|
cv2.rectangle(labeled, (x1, y1), (x2, y2), (0, 220, 0), 3)
|
|
cv2.putText(labeled, class_names[class_id], (x1, max(28, y1 - 7)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 220, 0), 2)
|
|
|
|
predicted = image.copy()
|
|
result = model.predict(str(image_path), imgsz=640, conf=threshold, device=0, verbose=False)[0]
|
|
predictions = []
|
|
for box in result.boxes:
|
|
class_id = int(box.cls.item())
|
|
confidence = float(box.conf.item())
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
|
|
predictions.append({"class": class_names[class_id], "confidence": round(confidence, 6), "xyxy": [x1, y1, x2, y2]})
|
|
cv2.rectangle(predicted, (x1, y1), (x2, y2), (0, 0, 255), 3)
|
|
cv2.putText(predicted, f"{class_names[class_id]} {confidence:.3f}", (x1, max(28, y1 - 7)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
|
|
|
|
canvas = cv2.hconcat([labeled, predicted])
|
|
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], 58), (255, 255, 255), -1)
|
|
cv2.putText(canvas, f"Original label: {image_path.name}", (20, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 120, 0), 2)
|
|
cv2.putText(canvas, "YOLO11s prediction", (width + 20, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 180), 2)
|
|
output_path = output_dir / f"{index:03d}_{image_path.stem}.jpg"
|
|
if not cv2.imwrite(str(output_path), canvas):
|
|
raise OSError(output_path)
|
|
manifest.append({"source": str(image_path), "label": str(label_path), "review_image": str(output_path), "predictions": predictions})
|
|
print(f"rendered {index}/{len(image_paths)} {image_path.name}", flush=True)
|
|
|
|
manifest_path = output_dir / "manifest.json"
|
|
import json
|
|
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(f"completed={len(manifest)} manifest={manifest_path}") |