track.py

"""
Run YOLOv8 + ByteTracker on a surgical video.

Each instrument gets a persistent ID across frames so you can follow
individual tools throughout the procedure.

Usage:
    python track.py --source surgery.mp4
    python track.py --source surgery.mp4 --weights runs/train/surgical/weights/best.pt
    python track.py --source 0           # webcam
"""

import argparse
import csv
from pathlib import Path

import cv2
from ultralytics import YOLO


CLASS_NAMES = [
    "Cautery",
    "Common bile duct",
    "Cystic Plate",
    "Cystic artery",
    "Cystic duct",
    "Duodenum",
    "Gallbladder",
    "Grasper",
    "Laparoscopic clip",
    "Laparoscopic clip applier",
    "Laparoscopic scissors",
    "Liver",
    "Omentum",
    "Plane of dissection",
    "Right colon",
    "Triangle of Calot",
]

# BGR colours — instruments get saturated colours, anatomy gets muted ones
PALETTE = [
    (0, 0, 255),    # Cautery              — red
    (180, 180, 0),  # Common bile duct     — dark cyan
    (0, 180, 180),  # Cystic Plate         — dark yellow
    (180, 0, 180),  # Cystic artery        — dark magenta
    (0, 130, 255),  # Cystic duct          — orange
    (100, 100, 100),# Duodenum             — grey
    (0, 200, 100),  # Gallbladder          — teal
    (0, 255, 0),    # Grasper              — bright green
    (255, 255, 0),  # Laparoscopic clip    — cyan
    (255, 200, 0),  # Laparoscopic clip applier — light blue
    (255, 0, 255),  # Laparoscopic scissors — magenta
    (150, 150, 150),# Liver                — light grey
    (80, 80, 80),   # Omentum              — dark grey
    (0, 100, 200),  # Plane of dissection  — amber
    (30, 30, 200),  # Right colon          — dark red
    (200, 100, 0),  # Triangle of Calot    — dark blue
]


def parse_args():
    p = argparse.ArgumentParser(description="Track surgical instruments in video")
    p.add_argument("--source", required=True,
                   help="Video file, image folder, or webcam index (0)")
    p.add_argument("--weights", default="runs/detect/surgical_s/weights/best.pt",
                   help="Path to trained YOLOv8 weights")
    p.add_argument("--conf", type=float, default=0.4,
                   help="Minimum detection confidence (0–1)")
    p.add_argument("--iou", type=float, default=0.5,
                   help="NMS IoU threshold")
    p.add_argument("--imgsz", type=int, default=640)
    p.add_argument("--save-video", action="store_true", default=True,
                   help="Save annotated output video")
    p.add_argument("--save-csv", action="store_true", default=False,
                   help="Save per-frame tracking CSV log")
    p.add_argument("--output-dir", default="results/tracking")
    p.add_argument("--show", action="store_true", default=False,
                   help="Display video while processing")
    p.add_argument("--device", default="cpu",
                   help="GPU index (0, 1, ...) or 'cpu'")
    return p.parse_args()


def annotate_frame(frame, boxes, track_ids, class_ids, confs):
    """Draw bounding boxes with track IDs directly on `frame` (in-place)."""
    for box, tid, cid, conf in zip(boxes, track_ids, class_ids, confs):
        x1, y1, x2, y2 = map(int, box)
        color = PALETTE[int(cid) % len(PALETTE)]
        label = f"ID{int(tid)} {CLASS_NAMES[int(cid)]} {conf:.0%}"

        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
        cv2.rectangle(frame, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
        cv2.putText(frame, label, (x1 + 2, y1 - 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1)
    return frame


def main():
    args = parse_args()

    weights_path = Path(args.weights)
    if not weights_path.exists():
        raise SystemExit(
            f"Weights not found: {weights_path}\n"
            "Train first with: python train.py\n"
            "Or pass --weights yolov8n.pt to use pretrained COCO weights (no surgical classes)."
        )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    model = YOLO(str(weights_path))

    # Use class names baked into the model weights if available
    global CLASS_NAMES, PALETTE
    if hasattr(model, "names") and model.names:
        model_names = [model.names[i] for i in sorted(model.names)]
        if model_names != CLASS_NAMES:
            CLASS_NAMES = model_names
            import random; random.seed(42)
            PALETTE = [(random.randint(50,255), random.randint(50,255), random.randint(50,255))
                       for _ in CLASS_NAMES]

    source_name = Path(str(args.source)).stem if str(args.source) != "0" else "webcam"
    out_video_path = output_dir / f"{source_name}_tracked.mp4"
    csv_path = output_dir / f"{source_name}_log.csv"

    csv_file = None
    csv_writer = None
    if args.save_csv:
        csv_file = open(csv_path, "w", newline="")
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(["frame", "track_id", "class", "conf", "x1", "y1", "x2", "y2"])

    # ByteTracker is built into ultralytics — `tracker="bytetrack.yaml"` enables it
    results_gen = model.track(
        source=args.source,
        conf=args.conf,
        iou=args.iou,
        imgsz=args.imgsz,
        device=args.device,
        tracker="bytetrack.yaml",
        stream=True,            # memory-efficient frame-by-frame generator
        persist=True,           # keep track state across frames
    )

    video_writer = None
    frame_idx = 0

    print(f"Processing {args.source} ...")
    for result in results_gen:
        frame = result.orig_img.copy()

        boxes = result.boxes
        if boxes is not None and boxes.id is not None:
            xyxy = boxes.xyxy.cpu().numpy()
            track_ids = boxes.id.cpu().numpy()
            class_ids = boxes.cls.cpu().numpy()
            confs = boxes.conf.cpu().numpy()

            annotate_frame(frame, xyxy, track_ids, class_ids, confs)

            if csv_writer:
                for box, tid, cid, conf in zip(xyxy, track_ids, class_ids, confs):
                    x1, y1, x2, y2 = map(int, box)
                    csv_writer.writerow([
                        frame_idx, int(tid), CLASS_NAMES[int(cid)],
                        f"{conf:.4f}", x1, y1, x2, y2,
                    ])

        if args.show:
            cv2.imshow("Surgical Instrument Tracking", frame)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break

        if args.save_video:
            if video_writer is None:
                h, w = frame.shape[:2]
                fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                video_writer = cv2.VideoWriter(str(out_video_path), fourcc, 25, (w, h))
            video_writer.write(frame)

        frame_idx += 1

    if video_writer:
        video_writer.release()
        print(f"Saved annotated video: {out_video_path}")

    if csv_file:
        csv_file.close()
        print(f"Saved tracking log:    {csv_path}")

    try:
        cv2.destroyAllWindows()
    except cv2.error:
        pass
    print(f"Processed {frame_idx} frames.")


if __name__ == "__main__":
    main()