Update app.py

app.py

@@ -21,16 +21,12 @@ def is_valid_image(file_content: bytes) -> Optional[str]:
 def process_image_bytes(image_bytes: bytes) -> np.ndarray:
     try:
         image = Image.open(io.BytesIO(image_bytes))
-        
-        if image.mode != 'RGB':
+        if image.mode == 'RGBA':
             image = image.convert('RGB')
-        
         img_array = np.array(image)
-        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
-            
         return img_array
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Error: {str(e)}")
 
 def get_mask(img, s=1024):
     img = (img / 255).astype(np.float32)
@@ -74,27 +70,25 @@ with gradio_app:
 
 @app.post("/remove-bg")
 async def remove_background(file: UploadFile = File(...)):
+    contents = await file.read()
+    
+    image_format = is_valid_image(contents)
+    if not image_format or image_format not in SUPPORTED_FORMATS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid format: {', '.join(SUPPORTED_FORMATS)}"
+        )
+    
     try:
-        contents = await file.read()
-        
-        image_format = is_valid_image(contents)
-        if not image_format or image_format not in SUPPORTED_FORMATS:
-            raise HTTPException(
-                status_code=400,
-                detail=f"Unsupported image format. Supported formats: {', '.join(SUPPORTED_FORMATS)}"
-            )
-        
         img = process_image_bytes(contents)
         mask = get_mask(img)
+        img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
+        mask = (mask * 255).astype(np.uint8)
+        img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
         
-        result = (mask * img).astype(np.uint8)
-        result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
-        alpha_channel = (mask * 255).astype(np.uint8)
-        rgba = np.dstack((result, alpha_channel))
-        
-        pil_image = Image.fromarray(rgba, 'RGBA')
+        pil_image = Image.fromarray(img, 'RGBA')
         img_byte_arr = io.BytesIO()
-        pil_image.save(img_byte_arr, format='PNG', optimize=True)
+        pil_image.save(img_byte_arr, format='PNG')
         img_byte_arr = img_byte_arr.getvalue()
         
         return Response(
@@ -105,10 +99,8 @@ async def remove_background(file: UploadFile = File(...)):
             }
         )
     
-    except HTTPException as he:
-        raise he
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
 
 if __name__ == "__main__":
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']