Spaces:
Sleeping
Sleeping
File size: 5,214 Bytes
e194ef7 | 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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | # ================================================================================
# YOLO Detection + Prediction Report
# ================================================================================
from ultralytics import YOLO
import os
# ββ EDIT THESE 4 PATHS ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
MODEL_PATH = r"D:\result\yolov8_microplastic\weights\best.pt"
TEST_IMAGES = r"D:\merged_dataset\test\images"
TEST_LABELS = r"D:\merged_dataset\test\labels"
OUTPUT_DIR = r"D:\result\inference_output"
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ STEP 1: Load model ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
print("Loading model...")
model = YOLO(MODEL_PATH)
print(f"Model loaded successfully. Total classes: {len(model.names)}")
# ββ STEP 2: Run detection on full test folder βββββββββββββββββββββββββββββββββββ
print(f"\nDetecting objects in: {TEST_IMAGES}")
print("Please wait...\n")
results = model.predict(
source = TEST_IMAGES,
conf = 0.25,
iou = 0.45,
imgsz = 640,
save = True,
save_txt = True,
save_conf = True,
project = OUTPUT_DIR,
name = "detections",
line_width = 2,
verbose = False,
)
print(f"Detection complete. Total images processed: {len(results)}")
# ββ STEP 3: Count correct / incorrect / missed ββββββββββββββββββββββββββββββββββ
correctly_predicted = []
incorrectly_predicted = []
no_detection = []
no_label = []
for result in results:
# get image filename without extension
stem = os.path.splitext(os.path.basename(result.path))[0]
lbl_path = os.path.join(TEST_LABELS, stem + ".txt")
# check if ground truth label exists
if not os.path.isfile(lbl_path):
no_label.append(stem)
continue
# read ground truth class IDs
with open(lbl_path, "r") as f:
lines = [l.strip() for l in f if l.strip()]
gt_classes = set(int(l.split()[0]) for l in lines)
# check if model detected anything
if result.boxes is None or len(result.boxes) == 0:
no_detection.append(stem)
continue
# get predicted class IDs
pred_classes = set(int(b.cls[0]) for b in result.boxes)
# correct = at least one predicted class matches ground truth
if gt_classes & pred_classes:
correctly_predicted.append(stem)
else:
incorrectly_predicted.append(stem)
# ββ STEP 4: Print report ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
total = len(results)
correct = len(correctly_predicted)
incorrect = len(incorrectly_predicted)
no_det = len(no_detection)
no_lbl = len(no_label)
print("\n" + "=" * 55)
print(" PREDICTION REPORT β Full Test Set")
print("=" * 55)
print(f" Total images : {total}")
print(f" Correctly predicted : {correct} ({correct/total*100:.1f}%)")
print(f" Incorrectly predicted : {incorrect} ({incorrect/total*100:.1f}%)")
print(f" No detection (missed) : {no_det} ({no_det/total*100:.1f}%)")
print(f" No label file found : {no_lbl}")
print("=" * 55)
# ββ STEP 5: Save report to file βββββββββββββββββββββββββββββββββββββββββββββββββ
os.makedirs(OUTPUT_DIR, exist_ok=True)
report_path = os.path.join(OUTPUT_DIR, "prediction_report.txt")
with open(report_path, "w") as f:
f.write("PREDICTION REPORT β Full Test Set\n")
f.write("=" * 55 + "\n")
f.write(f"Total images : {total}\n")
f.write(f"Correctly predicted : {correct} ({correct/total*100:.1f}%)\n")
f.write(f"Incorrectly predicted : {incorrect} ({incorrect/total*100:.1f}%)\n")
f.write(f"No detection (missed) : {no_det} ({no_det/total*100:.1f}%)\n")
f.write(f"No label file found : {no_lbl}\n")
f.write("\n\nCORRECTLY PREDICTED:\n")
for name in correctly_predicted:
f.write(f" {name}\n")
f.write("\n\nINCORRECTLY PREDICTED:\n")
for name in incorrectly_predicted:
f.write(f" {name}\n")
f.write("\n\nNO DETECTION (MISSED):\n")
for name in no_detection:
f.write(f" {name}\n")
f.write("\n\nNO LABEL FILE FOUND:\n")
for name in no_label:
f.write(f" {name}\n")
print(f"\n Annotated images saved : {os.path.join(OUTPUT_DIR, 'detections')}")
print(f" Full report saved : {report_path}") |