# ================================================================================
# 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}")