plots/plot.py

import cv2
import numpy as np
from pathlib import Path
from tqdm import tqdm
import matplotlib.pyplot as plt

from ultralytics import YOLO
from sklearn.metrics import (
    confusion_matrix, ConfusionMatrixDisplay,
    precision_recall_curve, auc,
    precision_recall_fscore_support
)
from sklearn.preprocessing import label_binarize
from torch.utils.tensorboard import SummaryWriter

# ─── CONFIG ────────────────────────────────────────────────────────────────────
MODEL_PATH     = "runs/detect/train/weights/best.pt"
VAL_IMG_DIR    = Path("C:/Users/rageb/Desktop/new yolo model/datasets/valid_filtered/images")
VAL_LABEL_DIR  = Path("C:/Users/rageb/Desktop/new yolo model/datasets/valid_filtered/labels")
CLASS_NAMES    = ['car', 'emv', 'htv']
IOU_THRESH     = 0.5
CONF_THRESH    = 0.25
TB_LOG_DIR     = "runs/eval"
# ────────────────────────────────────────────────────────────────────────────────

num_classes = len(CLASS_NAMES)
bg_idx = num_classes - 1  # index of “background”

def xywh2xyxy(xc, yc, w, h, img_w, img_h):
    """Convert YOLO normalized xc,yc,w,h → absolute x1,y1,x2,y2."""
    x1 = (xc - w/2) * img_w
    y1 = (yc - h/2) * img_h
    x2 = (xc + w/2) * img_w
    y2 = (yc + h/2) * img_h
    return [x1, y1, x2, y2]

def compute_iou(b1, b2):
    """Compute IoU of two [x1,y1,x2,y2] boxes."""
    xi1, yi1 = max(b1[0], b2[0]), max(b1[1], b2[1])
    xi2, yi2 = min(b1[2], b2[2]), min(b1[3], b2[3])
    inter_w, inter_h = max(0, xi2-xi1), max(0, yi2-yi1)
    inter = inter_w * inter_h
    area1 = (b1[2]-b1[0])*(b1[3]-b1[1])
    area2 = (b2[2]-b2[0])*(b2[3]-b2[1])
    union = area1 + area2 - inter
    return inter/union if union>0 else 0

# ─── Load model ────────────────────────────────────────────────────────────────
model = YOLO(MODEL_PATH)

# ─── Prepare holders ───────────────────────────────────────────────────────────
y_true = []
y_pred = []

# ─── Loop over validation images ───────────────────────────────────────────────
for img_path in tqdm(list(VAL_IMG_DIR.rglob("*.jpg")), desc="Evaluating"):
    img = cv2.imread(str(img_path))
    h, w = img.shape[:2]

    # 1) Run inference
    results = model(img, conf=CONF_THRESH)[0]
    pred_boxes, pred_cls, pred_scores = [], [], []
    if results.boxes is not None:
        for box in results.boxes:
            pred_boxes.append(box.xyxy.cpu().numpy()[0].tolist())
            pred_cls.append(int(box.cls.cpu().numpy()[0]))
            pred_scores.append(float(box.conf.cpu().numpy()[0]))

    # 2) Load ground‑truth boxes
    gt_file = VAL_LABEL_DIR / f"{img_path.stem}.txt"
    if not gt_file.exists():
        continue

    gt_boxes, gt_cls = [], []
    with open(gt_file) as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) < 5:
                # skip malformed
                continue
            cid = int(parts[0])
            xc, yc, ww, hh = map(float, parts[1:5])
            gt_cls.append(cid)
            gt_boxes.append(xywh2xyxy(xc, yc, ww, hh, w, h))

    # 3) Match predictions → GT by descending score & IoU
    matched_gt   = set()
    matched_pred = set()
    # pair (pb,pc,ps,pidx)
    preds = sorted(
        zip(pred_boxes, pred_cls, pred_scores, range(len(pred_boxes))),
        key=lambda x: x[2], reverse=True
    )
    for pb, pc, ps, pidx in preds:
        best_i, best_iou = -1, 0.0
        for gi, gb in enumerate(gt_boxes):
            if gi in matched_gt:
                continue
            iou = compute_iou(pb, gb)
            if iou > best_iou:
                best_iou, best_i = iou, gi
        if best_iou >= IOU_THRESH:
            # matched pair
            y_true.append(gt_cls[best_i])
            y_pred.append(pc)
            matched_gt.add(best_i)
            matched_pred.add(pidx)

    # 4) Unmatched predictions → false positives (background → pred_class)
    for pidx, pc in enumerate(pred_cls):
        if pidx not in matched_pred:
            y_true.append(bg_idx)
            y_pred.append(pc)

    # 5) Unmatched GT → false negatives (gt_class → background)
    for gi, gc in enumerate(gt_cls):
        if gi not in matched_gt:
            y_true.append(gc)
            y_pred.append(bg_idx)

# ─── Build & Save Confusion Matrix───────────────────────────────────────────
cm = confusion_matrix(y_true, y_pred, labels=range(num_classes))
cm_norm = cm.astype(float) / cm.sum(axis=1)[:,None]
disp = ConfusionMatrixDisplay(cm_norm, display_labels=CLASS_NAMES)
disp.plot(cmap=plt.cm.Blues)
plt.title("Normalized Confusion Matrix")
plt.savefig("confusion_matrix_normalized.png")
plt.close()

# ─── Precision–Recall Curves (per class) ───────────────────────────────────────
y_true_bin = label_binarize(y_true, classes=range(num_classes))
y_pred_bin = label_binarize(y_pred, classes=range(num_classes))

plt.figure()
for i in range(num_classes):
    prec, rec, _ = precision_recall_curve(y_true_bin[:,i], y_pred_bin[:,i])
    pr_auc = auc(rec, prec)
    plt.plot(rec, prec, label=f"{CLASS_NAMES[i]} (AUC={pr_auc:.2f})")
plt.xlabel("Recall"); plt.ylabel("Precision")
plt.title("Precision–Recall Curves")
plt.legend(loc="best")
plt.savefig("PR_curve.png")
plt.close()

# ─── Per‑class Precision / Recall / F1 ────────────────────────────────────────
prec, rec, f1, _ = precision_recall_fscore_support(
    y_true, y_pred, labels=range(num_classes)
)

plt.figure()
plt.plot(range(num_classes), prec, marker='o')
plt.xticks(range(num_classes), CLASS_NAMES)
plt.title("Precision per Class")
plt.savefig("P_curve.png")
plt.close()

plt.figure()
plt.plot(range(num_classes), rec, marker='o')
plt.xticks(range(num_classes), CLASS_NAMES)
plt.title("Recall per Class")
plt.savefig("R_curve.png")
plt.close()

plt.figure()
plt.plot(range(num_classes), f1, marker='o')
plt.xticks(range(num_classes), CLASS_NAMES)
plt.title("F1‑Score per Class")
plt.savefig("Fl_curve.png")
plt.close()

# ─── Label Correlation & Distribution ─────────────────────────────────────────
plt.figure()
plt.imshow(cm, cmap='hot', interpolation='nearest')
plt.colorbar()
plt.xticks(range(num_classes), CLASS_NAMES)
plt.yticks(range(num_classes), CLASS_NAMES)
plt.title("Label Correlation Matrix")
plt.savefig("labels_correlogram.jpg")
plt.close()

plt.figure()
counts = np.bincount(y_true, minlength=num_classes)
plt.bar(CLASS_NAMES, counts)
plt.title("Label Distribution")
plt.savefig("labels.jpg")
plt.close()

# ─── TensorBoard Logging ──────────────────────────────────────────────────────
writer = SummaryWriter(TB_LOG_DIR)
writer.add_scalar("Eval/Mean_Precision", np.mean(prec), 0)
writer.add_scalar("Eval/Mean_Recall",    np.mean(rec),  0)
writer.add_scalar("Eval/Mean_F1",        np.mean(f1),   0)
writer.close()

print("✅ Done! All plots saved and TensorBoard logs in", TB_LOG_DIR)