80 lines
3.4 KiB
Python
80 lines
3.4 KiB
Python
import csv
|
|
import json
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
from ultralytics import YOLO
|
|
|
|
root = Path(r"D:\work\yolo\data\Smoke-Fire-Detection-YOLO")
|
|
weights = Path(r"D:\work\yolo\runs\detect\smoke_fire_yolo11s_v1-4\weights\best.pt")
|
|
audit_dir = Path(r"D:\work\yolo\runs\audit")
|
|
csv_output = audit_dir / "suspected_missing_labels.csv"
|
|
json_output = audit_dir / "suspected_missing_labels.json"
|
|
image_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
|
|
|
|
images: list[str] = []
|
|
image_splits: dict[str, str] = {}
|
|
for split in ("train", "val", "test"):
|
|
for label_path in sorted((root / split / "labels").glob("*.txt")):
|
|
if label_path.read_text(encoding="utf-8", errors="replace").strip():
|
|
continue
|
|
for extension in image_extensions:
|
|
image_path = root / split / "images" / f"{label_path.stem}{extension}"
|
|
if image_path.is_file():
|
|
normalized = str(image_path.resolve())
|
|
images.append(normalized)
|
|
image_splits[normalized.lower()] = split
|
|
break
|
|
|
|
model = YOLO(str(weights))
|
|
rows: list[dict[str, object]] = []
|
|
chunk_size = 32
|
|
for offset in range(0, len(images), chunk_size):
|
|
chunk = images[offset : offset + chunk_size]
|
|
results = model.predict(source=chunk, imgsz=640, conf=0.25, device=0, batch=16, verbose=False)
|
|
for source_path, result in zip(chunk, results, strict=True):
|
|
boxes = result.boxes
|
|
if boxes is None or len(boxes) == 0:
|
|
continue
|
|
image_path = source_path
|
|
for class_id, confidence, xyxy in zip(boxes.cls.tolist(), boxes.conf.tolist(), boxes.xyxy.tolist()):
|
|
class_index = int(class_id)
|
|
rows.append({
|
|
"split": image_splits.get(image_path.lower(), "unknown"),
|
|
"image": image_path,
|
|
"class_id": class_index,
|
|
"class_name": model.names[class_index],
|
|
"confidence": round(float(confidence), 6),
|
|
"x1": round(float(xyxy[0]), 2),
|
|
"y1": round(float(xyxy[1]), 2),
|
|
"x2": round(float(xyxy[2]), 2),
|
|
"y2": round(float(xyxy[3]), 2),
|
|
})
|
|
if offset and offset % 1024 == 0:
|
|
print(f"scanned {offset}/{len(images)}", flush=True)
|
|
|
|
rows.sort(key=lambda row: float(row["confidence"]), reverse=True)
|
|
fieldnames = ["split", "image", "class_id", "class_name", "confidence", "x1", "y1", "x2", "y2"]
|
|
with csv_output.open("w", encoding="utf-8", newline="") as csv_file:
|
|
writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
|
|
thresholds = {}
|
|
for threshold in (0.25, 0.5, 0.7, 0.85):
|
|
selected = [row for row in rows if float(row["confidence"]) >= threshold]
|
|
thresholds[str(threshold)] = {
|
|
"detections": len(selected),
|
|
"images": len({row["image"] for row in selected}),
|
|
"smoke_images": len({row["image"] for row in selected if row["class_name"] == "smoke"}),
|
|
"fire_images": len({row["image"] for row in selected if row["class_name"] == "fire"}),
|
|
"by_split": dict(Counter(str(row["split"]) for row in selected)),
|
|
}
|
|
report = {
|
|
"empty_label_images": len(images),
|
|
"thresholds": thresholds,
|
|
"top_candidates": rows[:50],
|
|
"csv": str(csv_output),
|
|
}
|
|
json_output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps(report, ensure_ascii=False, indent=2)) |