Update app.py

app.py

@@ -1,15 +1,44 @@
+# app.py
 import gradio as gr
 import huggingface_hub
 import onnxruntime as rt
 import numpy as np
 import cv2
-from PIL import Image
-import io
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import Response
-import uvicorn
+import io
+from PIL import Image
+import imghdr
+from typing import Optional
 
-app = FastAPI()
+SUPPORTED_FORMATS = {'jpg', 'jpeg', 'png', 'bmp', 'webp', 'tiff'}
+
+def is_valid_image(file_content: bytes) -> Optional[str]:
+    image_format = imghdr.what(None, file_content)
+    if image_format is None:
+        return None
+    return image_format.lower()
+
+def process_image_bytes(image_bytes: bytes) -> np.ndarray:
+    try:
+        image = Image.open(io.BytesIO(image_bytes))
+        
+        if image.mode in ('RGBA', 'LA'):
+            background = Image.new('RGB', image.size, (255, 255, 255))
+            background.paste(image, mask=image.getchannel('A'))
+            image = background
+        
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        img_array = np.array(image)
+        
+        if len(img_array.shape) == 3:
+            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
+            
+        return img_array
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error: {str(e)}")
 
 def get_mask(img, s=1024):
     img = (img / 255).astype(np.float32)
@@ -26,11 +55,6 @@ def get_mask(img, s=1024):
     mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
     return mask
 
-def process_image(img_array):
-    mask = get_mask(img_array)
-    rgba = np.concatenate([img_array, (mask * 255).astype(np.uint8)], axis=2)
-    return rgba
-
 def rmbg_fn(img):
     mask = get_mask(img)
     img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
@@ -39,44 +63,64 @@ def rmbg_fn(img):
     mask = mask.repeat(3, axis=2)
     return mask, img
 
-@app.post("/remove-background")
+app = FastAPI()
+
+gradio_app = gr.Blocks()
+with gradio_app:
+    gr.Markdown("# Anime Remove Background\n\n"
+                "![visitor badge](https://api.visitorbadge.io/api/visitors?path=skytnt.animeseg&countColor=%23263759&style=flat&labelStyle=lower)\n\n"
+                "demo for [https://github.com/SkyTNT/anime-segmentation/](https://github.com/SkyTNT/anime-segmentation/)")
+    with gr.Column():
+        input_img = gr.Image(label="input image")
+        examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
+        examples = gr.Examples(examples=examples_data, inputs=[input_img])
+        run_btn = gr.Button(variant="primary")
+        with gr.Row():
+            output_mask = gr.Image(label="mask", format="png")
+            output_img = gr.Image(label="result", image_mode="RGBA", format="png")
+        run_btn.click(rmbg_fn, [input_img], [output_mask, output_img])
+
+@app.post("/remove-bg")
 async def remove_background(file: UploadFile = File(...)):
     contents = await file.read()
-    pil_image = Image.open(io.BytesIO(contents))
-    
-    if pil_image.mode != 'RGB':
-        pil_image = pil_image.convert('RGB')
-    img_array = np.array(pil_image)
     
-    result_array = process_image(img_array)
+    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)}"
+        )
     
-    result_image = Image.fromarray(result_array)
-    img_byte_arr = io.BytesIO()
-    result_image.save(img_byte_arr, format='PNG')
-    img_byte_arr = img_byte_arr.getvalue()
+    try:
+        img = process_image_bytes(contents)
+        
+        mask = get_mask(img)
+        result = (mask * img).astype(np.uint8)
+        
+        rgba = np.concatenate([result, (mask * 255).astype(np.uint8)], axis=2)
+        
+        pil_image = Image.fromarray(rgba)
+        img_byte_arr = io.BytesIO()
+        pil_image.save(img_byte_arr, format='PNG', optimize=True)
+        img_byte_arr = img_byte_arr.getvalue()
+        
+        return Response(
+            content=img_byte_arr, 
+            media_type="image/png",
+            headers={
+                "Content-Disposition": "attachment; filename=result.png"
+            }
+        )
     
-    return Response(content=img_byte_arr, media_type="image/png")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
 
 if __name__ == "__main__":
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
     model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
     rmbg_model = rt.InferenceSession(model_path, providers=providers)
     
-    gradio_app = gr.Blocks()
-    with gradio_app:
-        gr.Markdown("# Anime Remove Background\n\n"
-                   "![visitor badge](https://api.visitorbadge.io/api/visitors?path=skytnt.animeseg&countColor=%23263759&style=flat&labelStyle=lower)\n\n"
-                   "demo for [https://github.com/SkyTNT/anime-segmentation/](https://github.com/SkyTNT/anime-segmentation/)")
-        with gr.Column():
-            input_img = gr.Image(label="input image")
-            examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
-            examples = gr.Examples(examples=examples_data, inputs=[input_img])
-            run_btn = gr.Button(variant="primary")
-            with gr.Row():
-                output_mask = gr.Image(label="mask", format="png")
-                output_img = gr.Image(label="result", image_mode="RGBA", format="png")
-            run_btn.click(rmbg_fn, [input_img], [output_mask, output_img])
-    
-    gradio_app = gr.mount_gradio_app(app, gradio_app, path="/")
+    app = gr.mount_gradio_app(app, gradio_app, path="/")
     
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)