Upload 4 files

Dockerfile

@@ -18,6 +18,7 @@ RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r requirements.txt
 
 COPY --chown=user main.py .
+COPY --chown=user u2net.onnx .
 
 RUN mkdir -p uploads outputs
 

main.py

@@ -27,6 +27,12 @@ try:
 except ImportError:
     BG_AVAILABLE = False
 
+try:
+    from PIL import Image as PILImage
+    PIL_AVAILABLE = True
+except ImportError:
+    PIL_AVAILABLE = False
+
 basedir = os.path.abspath(os.path.dirname(__file__))
 load_dotenv(os.path.join(basedir, '.env'))
 
@@ -47,7 +53,9 @@ if BG_AVAILABLE:
         model_path = os.path.join(basedir, 'u2net.onnx')
         if os.path.exists(model_path):
             session = ort.InferenceSession(model_path)
-            logger.info("ONNX model loaded")
+            logger.info("ONNX u2net model loaded")
+        else:
+            logger.warning(f"u2net.onnx not found at {model_path}")
     except Exception as e:
         logger.warning(f"ONNX not loaded: {e}")
 
@@ -137,30 +145,39 @@ def index():
 def health():
     return jsonify({
         'status': 'healthy',
-        'background_removal': BG_AVAILABLE and session is not None
+        'background_removal': BG_AVAILABLE and session is not None,
+        'u2net_model': session is not None
     })
 
 
 @app.route('/remove', methods=['POST'])
 def remove():
     if not BG_AVAILABLE or session is None:
-        return jsonify({'error': 'Background removal not available'}), 503
+        return jsonify({'error': 'Background removal not available - u2net.onnx model missing'}), 503
     try:
         if 'image' not in request.files:
             return jsonify({'error': 'No image provided'}), 400
-        img = Image.open(request.files['image'].stream)
+        img = Image.open(request.files['image'].stream).convert('RGB')
         orig_size = img.size
-        inp = img.convert('RGB').resize((320, 320))
-        arr = np.expand_dims(np.transpose(np.array(inp).astype(np.float32) / 255.0, (2, 0, 1)), 0)
+        # Preprocess
+        resized = img.resize((320, 320))
+        arr = np.array(resized).astype(np.float32) / 255.0
+        arr = np.expand_dims(np.transpose(arr, (2, 0, 1)), 0)
+        # Run inference
         mask = session.run(None, {'input.1': arr})[0][0].squeeze()
-        mask = (mask - mask.min()) / (mask.max() - mask.min())
-        mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(orig_size)
+        # Normalize mask
+        mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
+        mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(orig_size, Image.LANCZOS)
+        # Apply mask as alpha channel
         out = img.convert('RGBA')
-        out.putdata([(r, g, b, m) for (r, g, b, a), m in zip(out.getdata(), mask.getdata())])
+        r_ch, g_ch, b_ch, _ = out.split()
+        out = Image.merge('RGBA', (r_ch, g_ch, b_ch, mask))
         buf = io.BytesIO()
         out.save(buf, format='PNG')
+        buf.seek(0)
         return jsonify({'image': 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode()})
     except Exception as e:
+        logger.error(f"Remove error: {e}", exc_info=True)
         return jsonify({'error': str(e)}), 500
 
 

requirements.txt

@@ -8,3 +8,6 @@ PyMuPDF==1.24.5
 python-docx==1.1.2
 openpyxl==3.1.5
 lxml==5.2.2
+pillow==10.4.0
+onnxruntime==1.18.1
+numpy==1.26.4