Skip to content

fix: YOLO model path, Docker CPU optimization, remove dead import #47

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kayceepeece wants to merge 1 commit into
base: main
Choose a base branch
from
+10 −3
Open

fix: YOLO model path, Docker CPU optimization, remove dead import #47

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import time
import cv2
import scenedetect
import subprocess
import argparse
import re
Expand Down Expand Up @@ -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
Expand Down