object-detect.py

from ultralytics import YOLO, RTDETR
from ultralytics.utils.metrics import box_iou
import torch
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn, retinanet_resnet50_fpn, ssdlite320_mobilenet_v3_large, fcos_resnet50_fpn
from torchvision.datasets import CocoDetection
from torchvision.transforms import functional as F
from torch.utils.data import DataLoader
from torchmetrics.detection.mean_ap import MeanAveragePrecision
from ultralytics.utils.metrics import box_iou
from PIL import Image
import cv2
import time as tm
import os


"""
Object Detection Testing Program

Tests various models using image and video test datasets, and evaluates accuracy
metrics such as precision, recall and Mean-Average-Precision.

Models that are evaulated:
1. YOLOv8
2. RT-DETR
3. Faster RCNN
4. RetinaNet
5. SSDLite
6. FCOS


Make sure to run with Python 3.12 venv as the interpreter
"""

DATASET = "val2017"

# SOURCE_TYPE: "dataset" = image dataset eval, "video" = video file, "mot17" = MOT17 frame sequence
SOURCE_TYPE = "dataset"
VIDEO_PATH = "/home/mispronounced/Code/Python Programs/Personal Projects/person-bicycle-car-detection.mp4"
VIDEO_GT_FILE = ""  # path to MOT17 gt.txt for video mode, or leave empty to skip accuracy metrics
MOT17_DIR = "/path/to/MOT17/train/MOT17-02-DPM"  # sequence dir containing img1/ and gt/gt.txt

if DATASET == "coco128":
    YOLO_DATA = "coco128.yaml"
    IMAGES_DIR = "/home/mispronounced/Code/JayRaj21/Object Detect/datasets/coco128/images/train2017"
    LABELS_DIR = "/home/mispronounced/Code/JayRaj21/Object Detect/datasets/coco128/labels/train2017"
    ANN_FILE = "/home/mispronounced/Code/JayRaj21/Object Detect/datasets/coco128/annotations.json"

elif DATASET == "val2017":
    YOLO_DATA = "coco.yaml"
    IMAGES_DIR = "/home/mispronounced/Code/JayRaj21/Object Detect/datasets/coco/images/val2017"
    LABELS_DIR = "/home/mispronounced/Code/JayRaj21/Object Detect/datasets/coco/labels/val2017"
    ANN_FILE = "/home/mispronounced/Code/JayRaj21/Object Detect/datasets/coco/annotations/instances_val2017.json"

def transform(img, target):
    return F.to_tensor(img), target

def collate_fn(x):
    return tuple(zip(*x))

dataset = CocoDetection(root=IMAGES_DIR, annFile=ANN_FILE, transforms=transform)

loader = DataLoader(dataset, batch_size=4, num_workers=2, collate_fn=collate_fn)
metric = MeanAveragePrecision(extended_summary=True)
metric.warn_on_many_detections = False

# model validation code
def validate(model):
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    time_old = tm.time()

    metric.reset()
    iou_scores = []

    # training process (skips gradient descent step)
    with torch.no_grad():
        for images, targets in loader:
            images = [img.to(device) for img in images]
            preds = model(images)

            formatted_targets = [{
                "boxes": torch.tensor([[a["bbox"][0], a["bbox"][1],
                                        a["bbox"][0]+a["bbox"][2],
                                        a["bbox"][1]+a["bbox"][3]] for a in t]).to(device),
                "labels": torch.tensor([a["category_id"] for a in t]).to(device)
            } for t in targets]

            metric.update(preds, formatted_targets)

            for pred, target in zip(preds, formatted_targets):
                if len(target['boxes']) == 0:
                    continue
                keep = pred['scores'] >= 0.3
                pred_boxes = pred['boxes'][keep].cpu()
                if len(pred_boxes) == 0:
                    continue
                iou_matrix = box_iou(pred_boxes, target['boxes'].cpu())
                iou_scores.append(iou_matrix.max(dim=1).values.mean().item())

    results = metric.compute()
    time_new = tm.time()

    p = results['precision']
    r = results['recall']
    print(f"mAP: {results['map']:.4f}")
    print(f"Precision: {p[p >= 0].mean():.4f}" if (p >= 0).any() else "Precision: N/A")
    print(f"Recall: {r[r >= 0].mean():.4f}" if (r >= 0).any() else "Recall: N/A")
    # print(f"Mean IoU: {sum(iou_scores)/len(iou_scores):.4f}" if iou_scores else "Mean IoU: N/A")
    print(f"Runtime: {time_new-time_old}")
    return

# execution code for running the models against the chosen dataset
def runModels(yolo_model, rcnn_model, ssdlite_model, fcos_model, rtdetr_model):
    for model_name, ul_model in [("YOLO", yolo_model), ("RT-DETR-L", rtdetr_model)]:
        print(f"\n--- {model_name} ---")
        time_old = tm.time()
        results = ul_model.val(data=YOLO_DATA)
        time_new = tm.time()
        print(f"mAP: {results.box.map}")
        print(f"Precision: {results.box.mp}")
        print(f"Recall: {results.box.mr}")
        print(f"Runtime: {time_new-time_old}")

        # iou_scores = []
        # for img_file in sorted(os.listdir(IMAGES_DIR)):
        #     if not img_file.endswith(".jpg"):
        #         continue
        #     img_path = os.path.join(IMAGES_DIR, img_file)
        #     label_path = os.path.join(LABELS_DIR, img_file.replace(".jpg", ".txt"))
        #     if not os.path.exists(label_path):
        #         continue
        #     from PIL import Image as PILImage
        #     img = PILImage.open(img_path)
        #     w, h = img.size
        #     gt_boxes = []
        #     with open(label_path) as f:
        #         for line in f:
        #             parts = line.strip().split()
        #             _, cx, cy, bw, bh = float(parts[0]), float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
        #             x1, y1 = (cx - bw/2) * w, (cy - bh/2) * h
        #             x2, y2 = (cx + bw/2) * w, (cy + bh/2) * h
        #             gt_boxes.append([x1, y1, x2, y2])
        #     if not gt_boxes:
        #         continue
        #     pred = ul_model.predict(img_path, verbose=False)[0]
        #     pred_boxes = pred.boxes.xyxy
        #     if len(pred_boxes) == 0:
        #         continue
        #     gt_tensor = torch.tensor(gt_boxes)
        #     iou_matrix = box_iou(pred_boxes.cpu(), gt_tensor)
        #     iou_scores.append(iou_matrix.max(dim=1).values.mean().item())
        # print(f"Mean IoU: {sum(iou_scores)/len(iou_scores):.4f}" if iou_scores else "Mean IoU: N/A")

    # Faster RCNN model
    validate(rcnn_model)

    # RetinaNet (ResNet50): accurate but slow on CPU
    # validate(retinanet_model)

    # SSDLite (MobileNetV3): faster on CPU, lower accuracy
    # validate(ssdlite_model)

    # FCOS (ResNet50): anchor-free single-stage detector
    validate(fcos_model)
    return

def load_mot_gt(gt_file):
    gt_by_frame = {}
    with open(gt_file) as f:
        for line in f:
            parts = line.strip().split(',')
            if len(parts) < 6:
                continue
            frame_id = int(parts[0])
            x, y, w, h = float(parts[2]), float(parts[3]), float(parts[4]), float(parts[5])
            gt_by_frame.setdefault(frame_id, []).append([x, y, x + w, y + h])
    return gt_by_frame

def runVideo(yolo_model, rcnn_model, fcos_model, rtdetr_model):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    gt_by_frame = load_mot_gt(VIDEO_GT_FILE) if VIDEO_GT_FILE and os.path.exists(VIDEO_GT_FILE) else {}
    has_gt = bool(gt_by_frame)

    torchvision_models = {"FasterRCNN": rcnn_model, "FCOS": fcos_model}

    # Ultralytics model video inference (YOLO and RT-DETR-S)
    for model_name, ul_model in [("YOLO", yolo_model), ("RT-DETR-L", rtdetr_model)]:
        print(f"\n--- {model_name} ---")
        time_old = tm.time()
        ul_results = ul_model.predict(VIDEO_PATH, save=True, verbose=False)
        print(f"Runtime: {tm.time() - time_old:.2f}s")
        if has_gt:
            iou_scores = []
            for frame_idx, result in enumerate(ul_results, start=1):
                if frame_idx not in gt_by_frame:
                    continue
                pred_boxes = result.boxes.xyxy.cpu()
                gt_tensor = torch.tensor(gt_by_frame[frame_idx])
                if len(pred_boxes) == 0:
                    continue
                iou_matrix = box_iou(pred_boxes, gt_tensor)
                iou_scores.append(iou_matrix.max(dim=1).values.mean().item())
            print(f"Mean IoU: {sum(iou_scores)/len(iou_scores):.4f}" if iou_scores else "Mean IoU: N/A")

    # torchvision model video inference
    for model_name, model in torchvision_models.items():
        print(f"\n--- {model_name} ---")
        model.eval()
        model.to(device)

        cap = cv2.VideoCapture(VIDEO_PATH)
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        out_path = f"output_{model_name}.mp4"
        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

        vid_metric = MeanAveragePrecision() if has_gt else None
        iou_scores = []
        frame_id = 1
        time_old = tm.time()

        with torch.no_grad():
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img_tensor = F.to_tensor(rgb).unsqueeze(0).to(device)
                pred = model(img_tensor)[0]

                keep = pred['scores'] >= 0.5
                pred_boxes = pred['boxes'][keep].cpu()
                pred_scores = pred['scores'][keep].cpu()
                pred_labels = pred['labels'][keep].cpu()

                if has_gt and frame_id in gt_by_frame:
                    gt_boxes = torch.tensor(gt_by_frame[frame_id])
                    gt_labels = torch.ones(len(gt_boxes), dtype=torch.long)
                    vid_metric.update(
                        [{"boxes": pred_boxes, "scores": pred_scores, "labels": pred_labels}],
                        [{"boxes": gt_boxes, "labels": gt_labels}]
                    )
                    if len(pred_boxes) > 0:
                        iou_matrix = box_iou(pred_boxes, gt_boxes)
                        iou_scores.append(iou_matrix.max(dim=1).values.mean().item())

                for box, score in zip(pred_boxes, pred_scores):
                    x1, y1, x2, y2 = box.int().tolist()
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f"{score:.2f}", (x1, max(y1 - 5, 0)),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

                writer.write(frame)
                frame_id += 1

        cap.release()
        writer.release()
        print(f"Runtime: {tm.time() - time_old:.2f}s  |  Output: {out_path}")

        if vid_metric:
            results = vid_metric.compute()
            print(f"mAP: {results['map']:.4f}")
            print(f"Mean IoU: {sum(iou_scores)/len(iou_scores):.4f}" if iou_scores else "Mean IoU: N/A")
        return

def runMOT17(yolo_model, rcnn_model, fcos_model, rtdetr_model):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    img_dir = os.path.join(MOT17_DIR, "img1")
    gt_file = os.path.join(MOT17_DIR, "gt", "gt.txt")
    gt_by_frame = load_mot_gt(gt_file) if os.path.exists(gt_file) else {}
    has_gt = bool(gt_by_frame)

    frame_files = sorted([f for f in os.listdir(img_dir) if f.endswith(".jpg")])

    # Ultralytics model inference (YOLO and RT-DETR-S)
    for model_name, ul_model in [("YOLO", yolo_model), ("RT-DETR-L", rtdetr_model)]:
        print(f"\n--- {model_name} ---")
        iou_scores = []
        time_old = tm.time()
        for frame_file in frame_files:
            frame_path = os.path.join(img_dir, frame_file)
            frame_id = int(os.path.splitext(frame_file)[0])
            result = ul_model.predict(frame_path, verbose=False)[0]
            if has_gt and frame_id in gt_by_frame:
                pred_boxes = result.boxes.xyxy.cpu()
                gt_tensor = torch.tensor(gt_by_frame[frame_id])
                if len(pred_boxes) > 0:
                    iou_matrix = box_iou(pred_boxes, gt_tensor)
                    iou_scores.append(iou_matrix.max(dim=1).values.mean().item())
        print(f"Runtime: {tm.time() - time_old:.2f}s")
        print(f"Mean IoU: {sum(iou_scores)/len(iou_scores):.4f}" if iou_scores else "Mean IoU: N/A")
    return

    # torchvision models
    torchvision_models = {"FasterRCNN": rcnn_model, "FCOS": fcos_model}
    first_frame = cv2.imread(os.path.join(img_dir, frame_files[0]))
    fh, fw = first_frame.shape[:2]

    for model_name, model in torchvision_models.items():
        print(f"\n--- {model_name} ---")
        model.eval()
        model.to(device)

        out_path = f"output_{model_name}.mp4"
        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (fw, fh))
        vid_metric = MeanAveragePrecision() if has_gt else None
        iou_scores = []
        time_old = tm.time()

        with torch.no_grad():
            for frame_file in frame_files:
                frame_path = os.path.join(img_dir, frame_file)
                frame_id = int(os.path.splitext(frame_file)[0])

                frame = cv2.imread(frame_path)
                rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img_tensor = F.to_tensor(rgb).unsqueeze(0).to(device)
                pred = model(img_tensor)[0]

                keep = pred['scores'] >= 0.5
                pred_boxes = pred['boxes'][keep].cpu()
                pred_scores = pred['scores'][keep].cpu()
                pred_labels = pred['labels'][keep].cpu()

                if has_gt and frame_id in gt_by_frame:
                    gt_boxes = torch.tensor(gt_by_frame[frame_id])
                    gt_labels = torch.ones(len(gt_boxes), dtype=torch.long)
                    vid_metric.update(
                        [{"boxes": pred_boxes, "scores": pred_scores, "labels": pred_labels}],
                        [{"boxes": gt_boxes, "labels": gt_labels}]
                    )
                    if len(pred_boxes) > 0:
                        iou_matrix = box_iou(pred_boxes, gt_boxes)
                        iou_scores.append(iou_matrix.max(dim=1).values.mean().item())

                for box, score in zip(pred_boxes, pred_scores):
                    x1, y1, x2, y2 = box.int().tolist()
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f"{score:.2f}", (x1, max(y1 - 5, 0)),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                writer.write(frame)

        writer.release()
        print(f"Runtime: {tm.time() - time_old:.2f}s  |  Output: {out_path}")
        if vid_metric:
            results = vid_metric.compute()
            print(f"mAP: {results['map']:.4f}")
            print(f"Mean IoU: {sum(iou_scores)/len(iou_scores):.4f}" if iou_scores else "Mean IoU: N/A")

# main code logic
if __name__ == "__main__":
    yolo_model = YOLO("yolov8m.pt")
    rtdetr_model = RTDETR("rtdetr-l.pt")
    rcnn_model = fasterrcnn_resnet50_fpn(pretrained=True)
    ssdlite_model = ssdlite320_mobilenet_v3_large(pretrained=True)
    fcos_model = fcos_resnet50_fpn(pretrained=True)

    if SOURCE_TYPE == "mot17":
        runMOT17(yolo_model, rcnn_model, fcos_model, rtdetr_model)
    elif SOURCE_TYPE == "video":
        runVideo(yolo_model, rcnn_model, fcos_model, rtdetr_model)
    else:
        for i in range(1):
            print(f"\n================== Run {i} ==================\n")
            runModels(yolo_model, rcnn_model, ssdlite_model, fcos_model, rtdetr_model)