phase2_video.py

"""
Phase 2b - Video Processing Pipeline
Downloads a sample warehouse video, runs YOLOv8 frame-by-frame,
saves annotated video output, and streams to Streamlit dashboard.

Run:
    # Download sample video
    python phase2_video.py --download

    # Process a video file
    python phase2_video.py --input data/sample_video/warehouse.mp4

    # Use webcam
    python phase2_video.py --webcam
"""

import argparse
import cv2
import time
import urllib.request
from pathlib import Path
from loguru import logger
from datetime import datetime


# ─── Config ───────────────────────────────────────────────────────
MODEL_NAME      = "yolov8n.pt"
CONF_THRESHOLD  = 0.35
OUTPUT_DIR      = Path("output/video")
VIDEO_DIR       = Path("data/sample_video")
FRAME_SKIP      = 2   # Process every Nth frame (speeds up CPU inference)

WAREHOUSE_LABELS = {
    "person":     ("worker",   (0, 200, 0)),
    "truck":      ("vehicle",  (0, 100, 255)),
    "car":        ("vehicle",  (0, 100, 255)),
    "motorcycle": ("vehicle",  (0, 100, 255)),
    "bicycle":    ("vehicle",  (0, 100, 255)),
    "forklift":   ("forklift", (0, 50, 255)),
    "suitcase":   ("parcel",   (255, 200, 0)),
    "backpack":   ("parcel",   (255, 200, 0)),
    "chair":      ("obstacle", (0, 0, 220)),
    "bottle":     ("item",     (200, 200, 200)),
    "handbag":    ("parcel",   (255, 200, 0)),
}

# ─── Sample warehouse videos (free, no auth needed) ───────────────
SAMPLE_VIDEOS = [
    {
        "name": "warehouse_workers.mp4",
        "url": "https://videos.pexels.com/video-files/3121461/3121461-uhd_2560_1440_25fps.mp4",
        "desc": "Warehouse workers and shelves"
    },
    {
        "name": "logistics_facility.mp4",
        "url": "https://videos.pexels.com/video-files/4570264/4570264-uhd_2560_1440_25fps.mp4",
        "desc": "Logistics facility operations"
    },
    {
        "name": "factory_floor.mp4",
        "url": "https://videos.pexels.com/video-files/3194277/3194277-uhd_2560_1440_25fps.mp4",
        "desc": "Factory floor with workers"
    },
]


def download_sample_video():
    """Download a sample warehouse video from Pexels (free, no API key)."""
    VIDEO_DIR.mkdir(parents=True, exist_ok=True)

    for item in SAMPLE_VIDEOS:
        dest = VIDEO_DIR / item["name"]
        if dest.exists():
            logger.info(f"Already exists: {item['name']}")
            return dest

        logger.info(f"Downloading: {item['name']} ({item['desc']})")
        logger.info("This may take 30-60 seconds depending on your connection...")

        try:
            def progress(count, block_size, total_size):
                if total_size > 0:
                    pct = count * block_size * 100 / total_size
                    print(f"\r  Progress: {min(pct, 100):.1f}%", end="", flush=True)

            urllib.request.urlretrieve(item["url"], dest, reporthook=progress)
            print()  # newline after progress
            logger.success(f"Downloaded: {dest}")
            return dest

        except Exception as e:
            logger.warning(f"Failed to download {item['name']}: {e}")
            continue

    logger.error("All downloads failed. Try uploading your own video via the dashboard.")
    return None


def load_model():
    """Load YOLOv8 model."""
    from ultralytics import YOLO
    logger.info(f"Loading YOLOv8 model: {MODEL_NAME}")
    model = YOLO(MODEL_NAME)
    logger.success("Model loaded ✓")
    return model


def draw_detections(frame, detections):
    """Draw bounding boxes on a frame."""
    for det in detections:
        x1, y1, x2, y2 = det["bbox"]
        colour = det["colour"]
        label  = f"{det['label']} {det['confidence']:.0%}"

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


def draw_stats(frame, frame_num, fps, total_detections):
    """Draw frame stats overlay."""
    overlay = f"Frame: {frame_num} | FPS: {fps:.1f} | Detections: {total_detections}"
    cv2.putText(frame, overlay, (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 3, cv2.LINE_AA)
    cv2.putText(frame, overlay, (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1, cv2.LINE_AA)
    return frame


def process_video(video_path: Path, show_live: bool = True, save_output: bool = True):
    """
    Process a video file with YOLOv8 detection.

    Args:
        video_path: Path to input video
        show_live:  Show live preview window
        save_output: Save annotated video to output/video/
    """
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        logger.error(f"Could not open video: {video_path}")
        return

    # Video properties
    fps      = cap.get(cv2.CAP_PROP_FPS) or 25
    width    = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    logger.info(f"Video: {video_path.name} | {width}x{height} | {fps:.1f}fps | {total_frames} frames")

    # Output video writer
    writer = None
    out_path = None
    if save_output:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        out_path  = OUTPUT_DIR / f"annotated_{video_path.stem}_{timestamp}.mp4"
        fourcc    = cv2.VideoWriter_fourcc(*"mp4v")
        writer    = cv2.VideoWriter(str(out_path), fourcc, fps, (width, height))
        logger.info(f"Saving to: {out_path}")

    # Load model
    model = load_model()

    frame_num       = 0
    total_detections = 0
    start_time      = time.time()

    logger.info("Processing... Press Q to quit live preview.")

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

        frame_num += 1

        # Skip frames for speed on CPU
        if frame_num % FRAME_SKIP != 0:
            if writer:
                writer.write(frame)
            continue

        # Run detection
        results = model(frame, conf=CONF_THRESHOLD, verbose=False)
        detections = []

        for result in results:
            for box in result.boxes:
                cls_name = result.names[int(box.cls)]
                info     = WAREHOUSE_LABELS.get(cls_name, (cls_name, (180, 180, 180)))
                wlabel, colour = info
                conf_val = float(box.conf)
                x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]
                detections.append({
                    "label":      wlabel,
                    "confidence": conf_val,
                    "bbox":       [x1, y1, x2, y2],
                    "colour":     colour,
                })

        total_detections += len(detections)

        # Draw annotations
        annotated = draw_detections(frame.copy(), detections)
        elapsed   = time.time() - start_time
        live_fps  = frame_num / elapsed if elapsed > 0 else 0
        annotated = draw_stats(annotated, frame_num, live_fps, len(detections))

        if writer:
            writer.write(annotated)

        if show_live:
            cv2.imshow("Warehouse Visual Intelligence - Press Q to quit", annotated)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                logger.info("Stopped by user")
                break

        # Progress log every 50 frames
        if frame_num % 50 == 0:
            pct = (frame_num / total_frames * 100) if total_frames > 0 else 0
            logger.info(f"  Frame {frame_num}/{total_frames} ({pct:.0f}%) | {live_fps:.1f} fps")

    cap.release()
    if writer:
        writer.release()
    cv2.destroyAllWindows()

    elapsed = time.time() - start_time
    logger.success(f"\n{'='*50}")
    logger.success(f"  Video processing complete!")
    logger.success(f"  Frames processed : {frame_num}")
    logger.success(f"  Total detections : {total_detections}")
    logger.success(f"  Time elapsed     : {elapsed:.1f}s")
    if out_path:
        logger.success(f"  Output saved     : {out_path}")
    logger.success(f"{'='*50}")

    return out_path


def process_webcam():
    """Run live YOLOv8 detection on webcam feed."""
    logger.info("Starting webcam... Press Q to quit.")
    model = load_model()
    cap   = cv2.VideoCapture(0)

    if not cap.isOpened():
        logger.error("Could not open webcam.")
        return

    frame_num  = 0
    start_time = time.time()

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

        frame_num += 1
        results    = model(frame, conf=CONF_THRESHOLD, verbose=False)
        detections = []

        for result in results:
            for box in result.boxes:
                cls_name = result.names[int(box.cls)]
                info     = WAREHOUSE_LABELS.get(cls_name, (cls_name, (180, 180, 180)))
                wlabel, colour = info
                conf_val = float(box.conf)
                x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]
                detections.append({
                    "label":      wlabel,
                    "confidence": conf_val,
                    "bbox":       [x1, y1, x2, y2],
                    "colour":     colour,
                })

        annotated = draw_detections(frame.copy(), detections)
        elapsed   = time.time() - start_time
        live_fps  = frame_num / elapsed if elapsed > 0 else 0
        annotated = draw_stats(annotated, frame_num, live_fps, len(detections))

        cv2.imshow("Warehouse Visual Intelligence - Webcam | Press Q to quit", annotated)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    cap.release()
    cv2.destroyAllWindows()
    logger.success("Webcam session ended.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Phase 2b: YOLOv8 Video Processing")
    parser.add_argument("--download", action="store_true", help="Download sample warehouse video")
    parser.add_argument("--input",    type=str, help="Path to video file")
    parser.add_argument("--webcam",   action="store_true", help="Use webcam")
    parser.add_argument("--no-save",  action="store_true", help="Don't save output video")
    parser.add_argument("--no-live",  action="store_true", help="Don't show live preview")
    parser.add_argument("--conf",     type=float, default=CONF_THRESHOLD, help="Confidence threshold")
    args = parser.parse_args()

    CONF_THRESHOLD = args.conf

    if args.download:
        video_path = download_sample_video()
        if video_path:
            logger.info(f"Video ready at: {video_path}")
            logger.info(f"Run: python phase2_video.py --input {video_path}")

    elif args.webcam:
        process_webcam()

    elif args.input:
        process_video(
            Path(args.input),
            show_live=not args.no_live,
            save_output=not args.no_save,
        )

    else:
        # Auto: download then process
        logger.info("No input specified — downloading sample video...")
        video_path = download_sample_video()
        if video_path:
            process_video(video_path, show_live=True, save_output=True)