Update Dockerfile and app.py for improved model handling and logging

Dockerfile

@@ -1,17 +1,39 @@
-FROM python:3.9
+FROM python:3.13-slim AS base
 
-# download this https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
-# copy model to avoid unnecessary download
-COPY u2net.onnx /home/.u2net/u2net.onnx
+# Use non-root user for better security
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    UV_SYSTEM_PIP=1
 
 WORKDIR /app
 
+# Install only essential system deps
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        libglib2.0-0 \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy uv binary for faster installs
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Pre-copy and install dependencies separately for layer caching
 COPY requirements.txt .
 
-RUN pip install --no-cache-dir -r requirements.txt
+RUN uv pip install --no-cache-dir -r requirements.txt --system && \
+    uv pip install --no-cache-dir gunicorn --system
 
+# Copy app source last to leverage Docker cache
 COPY . .
 
+# Pre-download u2net model to prevent runtime fetch
+RUN mkdir -p /root/.u2net && \
+    curl -L -o /root/.u2net/u2net.onnx \
+    https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
+
 EXPOSE 5100
 
-CMD ["python", "app.py"]
+# Use gunicorn instead of Flask dev server
+CMD ["gunicorn", "--bind", "0.0.0.0:5100", "--workers", "2", "app:app"]

app.py

@@ -2,26 +2,73 @@
 from rembg import remove
 from PIL import Image
 from io import BytesIO
+import os
+import logging
+import urllib.request
+from time import perf_counter
 
 app = Flask(__name__)
 
+# Configure structured logging early
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s | %(levelname)s | %(name)s | %(message)s'
+)
+logger = logging.getLogger("rmbg-app")
+
+MODEL_PATH = os.path.expanduser('~/.u2net/u2net.onnx')
+MODEL_URL = 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx'
+
+
+def ensure_model_cached():
+    """Ensure model exists locally; download once."""
+    if os.path.exists(MODEL_PATH):
+        return
+    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+    logger.info("Downloading u2net.onnx model...")
+    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
+    logger.info(f"Model cached at {MODEL_PATH}")
+
+
+@app.before_request
+def preload_model():
+    """Run model check once before first request."""
+    ensure_model_cached()
+    logger.info("Model ready for inference.")
+
+
 @app.route('/', methods=['GET', 'POST'])
 def upload_file():
     if request.method == 'POST':
-        if 'file' not in request.files:
+        start = perf_counter()
+
+        file = request.files.get('file')
+        if not file or file.filename == '':
+            logger.warning("No file provided in request.")
             return 'No file uploaded', 400
-        file = request.files['file']
-        if file.filename == '':
-            return 'No file selected', 400
-        if file:
+
+        try:
             input_image = Image.open(file.stream)
             output_image = remove(input_image, post_process_mask=True)
-            img_io = BytesIO()
-            output_image.save(img_io, 'PNG')
-            img_io.seek(0)
-            # return send_file(img_io, mimetype='image/png')  # Change download in separatre browser tab
-            return send_file(img_io, mimetype='image/png', as_attachment=True, download_name='_rmbg.png')
+        except Exception as e:
+            logger.exception(f"Failed to process image: {e}")
+            return 'Error processing image', 500
+
+        img_io = BytesIO()
+        output_image.save(img_io, 'PNG')
+        img_io.seek(0)
+
+        elapsed = perf_counter() - start
+        logger.info(f"Processed {file.filename} in {elapsed:.3f}s")
+
+        return send_file(img_io, mimetype='image/png',
+                         as_attachment=True, download_name='_rmbg.png')
+
     return render_template('index.html')
 
+
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', debug=True, port=5100)
+    # Preload model for faster cold start when container spins up
+    ensure_model_cached()
+    logger.info("Starting Flask server on port 5100")
+    app.run(host='0.0.0.0', port=5100, debug=False)

requirements.txt

@@ -1,3 +1,4 @@
 flask
 rembg
-pillow
+pillow
+onnxruntime