""" 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)