From f5b8a05518605365c22c20111e42df8ad06916c3 Mon Sep 17 00:00:00 2001
From: Kaycee <kayceepeece@users.noreply.github.com>
Date: Wed, 24 Jun 2026 17:23:41 +0000
Subject: [PATCH] fix: YOLO model path, Docker CPU optimization, remove dead
 import
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

1. Make YOLO model path configurable via YOLO_MODEL_PATH env var
   - Default: yolov8n.pt (auto-downloads to CWD for bare-metal)
   - Docker: /tmp/Ultralytics/yolov8n.pt (pre-downloaded, survives volume mount)
   - Fixes PermissionError when docker-compose mounts .:/app

2. CPU-only PyTorch by default in Docker
   - Split pip install: torch via --index-url, rest via requirements.txt
   - TORCH_INDEX_URL build arg (default: cpu)
   - GPU users: --build-arg TORCH_INDEX_URL=https://download.pytorch.org/whl/cu124
   - Image: 11.3 GB → 4.6 GB

3. Remove unused import scenedetect (dead code after open_video refactor)
---
 Dockerfile | 9 ++++++++-
 main.py    | 4 ++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 95bfd1f8..b81c8814 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -13,7 +13,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 COPY requirements.txt .
 RUN python -m venv /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
+# Install PyTorch CPU-only first (smaller image: ~4.6 GB vs ~11 GB).
+# GPU users: docker compose build --build-arg TORCH_INDEX_URL=https://download.pytorch.org/whl/cu124
+ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu
 RUN pip install --upgrade pip
+RUN pip install --no-cache-dir --index-url ${TORCH_INDEX_URL} torch==2.11.0 torchvision==0.26.0
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Final stage
@@ -37,6 +41,9 @@ COPY --from=builder /opt/venv /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
 ENV PYTHONUNBUFFERED=1
 
+# Point YOLO to pre-downloaded model (survives volume mount)
+ENV YOLO_MODEL_PATH=/tmp/Ultralytics/yolov8n.pt
+
 # Always upgrade yt-dlp to latest (YouTube bot-detection changes frequently)
 RUN pip install --upgrade --no-cache-dir yt-dlp
 
@@ -55,7 +62,7 @@ RUN chown -R appuser:appuser /app /tmp/Ultralytics
 USER appuser
 
 # Pre-download YOLO model on build (now running as appuser)
-RUN python -c "from ultralytics import YOLO; YOLO('yolov8n.pt')"
+RUN python -c "from ultralytics import YOLO; YOLO('/tmp/Ultralytics/yolov8n.pt')"
 
 # Expose FastAPI port
 EXPOSE 8000
diff --git a/main.py b/main.py
index 0e26f8ce..f594deaf 100644
--- a/main.py
+++ b/main.py
@@ -1,6 +1,5 @@
 import time
 import cv2
-import scenedetect
 import subprocess
 import argparse
 import re
@@ -68,7 +67,8 @@
 """
 
 # Load the YOLO model once (Keep for backup or scene analysis if needed)
-model = YOLO('yolov8n.pt')
+YOLO_MODEL_PATH = os.environ.get("YOLO_MODEL_PATH", "yolov8n.pt")
+model = YOLO(YOLO_MODEL_PATH)
 
 # --- MediaPipe Setup ---
 # Use standard Face Detection (BlazeFace) for speed