Spaces:
Sleeping
Sleeping
File size: 4,006 Bytes
cabc6bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Hunt label errors in the dataset with the trained model + FiftyOne (plan: data audit).
Where the model confidently disagrees with the ground truth, the *label* is
often what's wrong — especially for the NO-* violation classes, where missed
annotations are easiest to make. This script loads a split, attaches model
predictions, computes FiftyOne's mistakenness score, and opens the FiftyOne app
sorted so the most suspicious annotations come first.
Needs a local dataset copy (scripts/download_data.py) and the `audit` group:
uv run --group audit python scripts/audit_labels.py # train split, opens the app
uv run --group audit python scripts/audit_labels.py --split valid
uv run --group audit python scripts/audit_labels.py --no-app --top 50
Review workflow: walk the sorted samples in the app, fix bad boxes/classes in
your labeling tool (Roboflow/CVAT) — FiftyOne is the *finder*, not the editor.
Expect a meaningful hit-rate in the top few hundred; stop when flags stop
being real errors.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import fiftyone as fo
import fiftyone.brain as fob
from ultralytics import YOLO
WEIGHTS = "runs/detect/yolov8n_v1_train/weights/best.pt"
DATA_DIR = "work"
def attach_predictions(dataset: fo.Dataset, weights: str, conf: float, imgsz: int) -> None:
model = YOLO(weights)
names = model.names
with fo.ProgressBar() as pb:
for sample in pb(dataset):
r = model.predict(sample.filepath, conf=conf, imgsz=imgsz, verbose=False)[0]
dets = []
for box in r.boxes:
x1, y1, x2, y2 = box.xyxyn[0].tolist()
dets.append(
fo.Detection(
label=names[int(box.cls)],
bounding_box=[x1, y1, x2 - x1, y2 - y1],
confidence=float(box.conf),
)
)
sample["predictions"] = fo.Detections(detections=dets)
sample.save()
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--data-dir", default=DATA_DIR)
p.add_argument("--weights", default=WEIGHTS)
p.add_argument("--split", choices=["train", "valid", "test"], default="train")
p.add_argument("--conf", type=float, default=0.15, help="low on purpose — misses need candidates to compare")
p.add_argument("--imgsz", type=int, default=640)
p.add_argument("--top", type=int, default=25, help="how many most-suspicious samples to print")
p.add_argument("--no-app", action="store_true", help="print the ranking only, don't launch the FiftyOne app")
args = p.parse_args()
if not (Path(args.data_dir) / args.split / "images").is_dir():
raise SystemExit(f"{args.data_dir}/{args.split}/images not found — fetch with scripts/download_data.py")
name = f"ppe-audit-{args.split}"
if fo.dataset_exists(name):
fo.delete_dataset(name) # fresh predictions each run beat stale cached ones
dataset = fo.Dataset.from_dir(
name=name,
dataset_type=fo.types.YOLOv5Dataset,
dataset_dir=args.data_dir,
yaml_path="data.yaml",
split=args.split,
label_field="ground_truth",
)
print(f"Loaded {len(dataset)} {args.split} samples")
print("Attaching model predictions ...")
attach_predictions(dataset, args.weights, args.conf, args.imgsz)
print("Computing mistakenness ...")
fob.compute_mistakenness(dataset, "predictions", label_field="ground_truth")
view = dataset.sort_by("mistakenness", reverse=True)
print(f"\nTop {args.top} most-suspicious samples (highest mistakenness first):")
for sample in view[: args.top]:
print(f" {sample.mistakenness:.3f} {Path(sample.filepath).name}")
if not args.no_app:
session = fo.launch_app(view)
session.wait() # keep the app alive until the user closes it
if __name__ == "__main__":
main()
|