diff --git a/Dockerfile b/Dockerfile index 0182438..1958a30 100644 --- a/Dockerfile +++ b/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"] \ No newline at end of file +# Use gunicorn instead of Flask dev server +CMD ["gunicorn", "--bind", "0.0.0.0:5100", "--workers", "2", "app:app"] diff --git a/app.py b/app.py index c753d15..353845b 100644 --- a/app.py +++ b/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) \ No newline at end of file + # 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) diff --git a/requirements.txt b/requirements.txt index 10ac15b..a7cf01c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ flask rembg -pillow \ No newline at end of file +pillow +onnxruntime \ No newline at end of file