feat: enhance Docker setup with entrypoint script and model downloads

Author: tuke307

tree-disk-api/.dockerignore

@@ -33,3 +33,4 @@ docker.compose.yml
 output/
 input/
 tests/
+models/

tree-disk-api/Dockerfile

@@ -38,9 +38,16 @@ COPY --from=builder /app/.venv .venv
 # Copy application code
 COPY . .
 
+# Add the source directory to PYTHONPATH so imports resolve correctly
+ENV PYTHONPATH="/app/src"
+
+# Copy the entrypoint script into the image and set execute permission
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
 # Activate virtual environment
 ENV PATH="/app/.venv/bin:$PATH"
 
 EXPOSE 3100
 
-CMD ["gunicorn", "main:app", "-c", "gunicorn.conf.py"]
+ENTRYPOINT ["/entrypoint.sh"]

tree-disk-api/entrypoint.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+set -e
+
+echo "Downloading models..."
+python src/config/preparation.py
+
+echo "Starting server..."
+exec gunicorn main:app -c gunicorn.conf.py

tree-disk-api/poetry.lock

@@ -16,6 +16,7 @@ python-multipart = "^0.0.17"
 tree-disk-segmentation = "^0.1.3"
 tree-disk-pith = "^0.1.4"
 tree-disk-rings = "^0.4.5" 
+gdown = "^5.2.0"
 
 [poetry.group.dev.dependencies]
 tree-disk-segmentation = { path = "../tree-disk-segmentation" }

tree-disk-api/pyproject.toml

@@ -4,7 +4,7 @@
 from PIL import Image
 import io
 
-from config import OUTPUT_DIR, INPUT_DIR, YOLO_MODEL_PATH, DEBUG, SAVE_RESULTS
+from config.settings import OUTPUT_DIR, INPUT_DIR, YOLO_MODEL_PATH, DEBUG, SAVE_RESULTS
 
 router = APIRouter()
 

tree-disk-api/src/api/endpoints/pith.py

@@ -4,7 +4,7 @@
 import io
 from io import BytesIO
 
-from config import OUTPUT_DIR, INPUT_DIR, DEBUG, SAVE_RESULTS
+from config.settings import OUTPUT_DIR, INPUT_DIR, DEBUG, SAVE_RESULTS
 
 router = APIRouter()
 

tree-disk-api/src/api/endpoints/rings.py

@@ -4,7 +4,7 @@
 import io
 from io import BytesIO
 
-from config import OUTPUT_DIR, INPUT_DIR, U2NET_MODEL_PATH, DEBUG, SAVE_RESULTS
+from config.settings import OUTPUT_DIR, INPUT_DIR, U2NET_MODEL_PATH, DEBUG, SAVE_RESULTS
 
 router = APIRouter()
 

tree-disk-api/src/api/endpoints/segmentation.py

@@ -0,0 +1,49 @@
+import logging
+import gdown
+from config.settings import (
+    INPUT_DIR,
+    OUTPUT_DIR,
+    MODEL_PATH,
+    U2NET_MODEL_PATH,
+    YOLO_MODEL_PATH,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def download_u2net_model():
+    """
+    Download the model from Google Drive.
+    """
+    file_id = "1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ"
+    url = f"https://drive.google.com/uc?id={file_id}"
+
+    logger.info(f"Downloading Google Drive file from: {url}")
+
+    gdown.download(url, str(U2NET_MODEL_PATH), quiet=False)
+
+
+def download_yolo_model():
+    """
+    Download the model from Google Drive.
+    """
+    file_id = "1_-dDH4DNiL8wbgiPWPaNqne7kbJyZc8z"
+    url = f"https://drive.google.com/uc?id={file_id}"
+
+    logger.info(f"Downloading Google Drive file from: {url}")
+
+    gdown.download(url, str(YOLO_MODEL_PATH), quiet=False)
+
+
+def create_dirs():
+    """
+    Create necessary directories if they do not exist.
+    """
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    INPUT_DIR.mkdir(parents=True, exist_ok=True)
+    MODEL_PATH.mkdir(parents=True, exist_ok=True)
+
+
+create_dirs()
+download_u2net_model()
+download_yolo_model()