Merge pull request #4 from Roalkege/main

This PR improves the out-of-the-box experience for running `light-object-detect`, especially in Docker / on Unraid.

Author: matteius

Dockerfile

@@ -34,6 +34,11 @@ RUN python -m pip install --upgrade pip \
 
 COPY . .
 
+# Bake the default TFLite model into the image (can be disabled at build time).
+# Usage: docker build --build-arg DOWNLOAD_DEFAULT_MODEL=0 ...
+ARG DOWNLOAD_DEFAULT_MODEL=1
+RUN if [ "$DOWNLOAD_DEFAULT_MODEL" = "1" ]; then python scripts/download_model.py; fi
+
 EXPOSE 8000
 
 # .env is optional; Unraid can mount it to /app/.env

README.md

@@ -28,11 +28,13 @@ A lightweight Python API for object detection with pluggable backends. This API
    pipenv install
    ```
 
-3. Download a sample TFLite model:
+3. (Optional) Pre-download the default TFLite model:
    ```bash
    pipenv run python scripts/download_model.py
    ```
 
+   If you skip this step, the API will try to download the default model on startup when the `tflite` backend is enabled (requires internet access). Docker builds download the default model by default.
+
 ## Usage
 
 1. Start the API server using the provided script:
@@ -49,23 +51,26 @@ A lightweight Python API for object detection with pluggable backends. This API
 
 3. Access the API documentation at http://localhost:9001/docs
 
-## Docker (z.B. Unraid / lightNVR)
+## Docker (e.g. Unraid / lightNVR)
 
 ### Build
 
 ```bash
 docker build -t light-object-detect:local .
 ```
 
+By default, the image downloads a small reference TFLite model at build time so the `tflite` backend works out of the box.
+To disable this, build with `--build-arg DOWNLOAD_DEFAULT_MODEL=0`.
+
 ### Run
 
-Option A: ohne `.env` (Defaults aus `config.py`):
+Option A: without `.env` (uses defaults from `config.py`):
 
 ```bash
 docker run --rm -p 8000:8000 --name light-object-detect light-object-detect:local
 ```
 
-Option B: mit `.env` (empfohlen, z.B. Backend/Model-Pfade):
+Option B: with `.env` (recommended, e.g. for backend/model paths):
 
 ```bash
 docker run --rm -p 8000:8000 --name light-object-detect \
@@ -86,7 +91,7 @@ docker run --rm -p 8000:8000 --name light-object-detect `
 
 ### lightNVR Integration
 
-In lightNVR als API-URL typischerweise:
+In lightNVR, the API URL is typically:
 
 - `http://<docker-host>:8000/api/v1/detect`
 

main.py

@@ -7,6 +7,7 @@
 from api.router import router as api_router
 from config import settings
 from backends.factory import get_backend
+from utils.model_download import ensure_tflite_ssd_mobilenet_v1, ModelDownloadError
 
 # Configure logging
 logging.basicConfig(
@@ -30,6 +31,25 @@ async def startup_event():
     logger.info(f"Available backends: {settings.AVAILABLE_BACKENDS}")
     logger.info(f"Default backend: {settings.DEFAULT_BACKEND}")
 
+    # Ensure the default TFLite model is present if TFLite is enabled.
+    # This keeps the out-of-the-box experience working, while still allowing
+    # users to override paths via .env.
+    if "tflite" in settings.AVAILABLE_BACKENDS:
+        try:
+            ensure_result = ensure_tflite_ssd_mobilenet_v1(
+                model_path=settings.TFLITE_MODEL_PATH,
+                labels_path=settings.TFLITE_LABELS_PATH,
+                force=False,
+            )
+            if ensure_result.did_download:
+                logger.info(f"✓ Downloaded default TFLite model: {ensure_result.message}")
+        except ModelDownloadError as e:
+            logger.error(
+                "✗ TFLite model is missing and could not be downloaded. "
+                f"Reason: {e}. "
+                "Fix: run 'python scripts/download_model.py' or set TFLITE_MODEL_PATH to an existing model."
+            )
+
     # Test each backend
     for backend_name in settings.AVAILABLE_BACKENDS:
         try:

scripts/download_model.py

@@ -1,103 +1,39 @@
 #!/usr/bin/env python3
 """
-Script to download a sample TFLite model for object detection.
+Script to download a default TFLite model for object detection.
+
+This is mainly intended for local development; Docker builds can also run it
+to bake the model into the image.
 """
 
-import os
+from __future__ import annotations
+
 import sys
-import urllib.request
-import zipfile
-import shutil
 from pathlib import Path
 
-# Add parent directory to path to import config
+# Add parent directory to path to import config + utils
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-from config import settings
-
-
-def download_file(url, destination):
-    """Download a file from a URL to a destination."""
-    print(f"Downloading {url} to {destination}...")
-    
-    # Create directory if it doesn't exist
-    os.makedirs(os.path.dirname(destination), exist_ok=True)
-    
-    # Download the file
-    urllib.request.urlretrieve(url, destination)
-    
-    print(f"Downloaded {destination}")
 
-
-def extract_zip(zip_path, extract_dir):
-    """Extract a zip file to a directory."""
-    print(f"Extracting {zip_path} to {extract_dir}...")
-    
-    # Create directory if it doesn't exist
-    os.makedirs(extract_dir, exist_ok=True)
-    
-    # Extract the zip file
-    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        zip_ref.extractall(extract_dir)
-    
-    print(f"Extracted {zip_path}")
-
-
-def download_ssd_mobilenet():
-    """Download SSD MobileNet v1 model."""
-    model_url = "https://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip"
-    zip_path = "/tmp/ssd_mobilenet.zip"
-    extract_dir = "/tmp/ssd_mobilenet"
-    model_dir = os.path.dirname(settings.TFLITE_MODEL_PATH)
-    
-    # Download the model
-    download_file(model_url, zip_path)
-    
-    # Extract the zip file
-    extract_zip(zip_path, extract_dir)
-    
-    # Create model directory if it doesn't exist
-    os.makedirs(model_dir, exist_ok=True)
-    
-    # Copy the model file
-    shutil.copy(
-        os.path.join(extract_dir, "detect.tflite"),
-        settings.TFLITE_MODEL_PATH
-    )
-    
-    # Create a labelmap file
-    with open(settings.TFLITE_LABELS_PATH, 'w') as f:
-        # COCO dataset labels
-        labels = [
-            "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
-            "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
-            "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
-            "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
-            "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
-            "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
-            "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
-            "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
-            "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book",
-            "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
-        ]
-        for label in labels:
-            f.write(f"{label}\n")
-    
-    # Clean up
-    os.remove(zip_path)
-    shutil.rmtree(extract_dir)
-    
-    print(f"Model saved to {settings.TFLITE_MODEL_PATH}")
-    print(f"Labels saved to {settings.TFLITE_LABELS_PATH}")
+from config import settings
+from utils.model_download import ensure_tflite_ssd_mobilenet_v1, ModelDownloadError
 
 
 def main():
     """Main function."""
-    print("Downloading TFLite model for object detection...")
-    
-    # Download SSD MobileNet model
-    download_ssd_mobilenet()
-    
-    print("Done!")
+    print("Ensuring default TFLite model for object detection...")
+
+    try:
+        result = ensure_tflite_ssd_mobilenet_v1(
+            model_path=settings.TFLITE_MODEL_PATH,
+            labels_path=settings.TFLITE_LABELS_PATH,
+            force=False,
+        )
+    except ModelDownloadError as e:
+        print(f"ERROR: {e}")
+        raise SystemExit(1) from e
+
+    print(result.message)
+    print("Done.")
 
 
 if __name__ == "__main__":