Update app.py

app.py

@@ -1,113 +1,50 @@
-# app.py
-import gradio as gr
-import torch
-import numpy as np
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import Response
 from PIL import Image
 import io
-import base64
-import requests
-import os
-from torchvision.transforms import Compose, Resize, ToTensor, Normalize
-
-# Download pre-trained DIS (IS-Net) weights
-def download_weights():
-    weights_path = "isnet-general-use.pth"
-    if not os.path.exists(weights_path):
-        url = "https://github.com/xuebinqin/DIS/releases/download/v1.0/isnet-general-use.pth"
-        try:
-            response = requests.get(url, stream=True)
-            response.raise_for_status()
-            with open(weights_path, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
-        except Exception as e:
-            raise Exception(f"Failed to download weights: {str(e)}")
-    return weights_path
-
-# DIS (IS-Net) model architecture (simplified from https://github.com/xuebinqin/DIS)
-class ISNet(torch.nn.Module):
-    def __init__(self):
-        super(ISNet, self).__init__()
-        # Simplified architecture (for demonstration; replace with full IS-Net)
-        # Full architecture: https://github.com/xuebinqin/DIS/blob/main/ISNet.py
-        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1)
-        self.pool = torch.nn.MaxPool2d(2, 2)
-        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1)
-        self.upconv = torch.nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
-        self.conv3 = torch.nn.Conv2d(64, 1, kernel_size=3, padding=1)
+import numpy as np
+from transformers import AutoModelForImageSegmentation, AutoProcessor
+import torch
 
-    def forward(self, x):
-        # Simplified forward pass (replace with full IS-Net forward)
-        x = torch.relu(self.conv1(x))
-        x = self.pool(x)
-        x = torch.relu(self.conv2(x))
-        x = self.upconv(x)
-        x = torch.sigmoid(self.conv3(x))
-        return x
+app = FastAPI()
 
-# Initialize model
-try:
-    weights_path = download_weights()
-    model = ISNet()
-    state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
-    model.load_state_dict(state_dict)
-    model.eval()
-except Exception as e:
-    raise Exception(f"Model initialization failed: {str(e)}")
+# Load the RMBG V1.4 model and processor
+model = AutoModelForImageSegmentation.from_pretrained(
+    "briaai/RMBG-1.4", trust_remote_code=True
+)
+processor = AutoProcessor.from_pretrained("briaai/RMBG-1.4")
 
-def remove_background(image):
-    """
-    Remove background using DIS (IS-Net).
-    Input: PIL Image
-    Output: Base64-encoded PNG with transparent background
-    """
+@app.post("/remove-background")
+async def remove_background(file: UploadFile = File(...)):
     try:
-        # Ensure image is RGB
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        
+        # Read uploaded image
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+
         # Preprocess image
-        transform = Compose([
-            Resize((1024, 1024)),
-            ToTensor(),
-            Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-        ])
-        img_tensor = transform(image).unsqueeze(0)
+        inputs = processor(images=image, return_tensors="pt")
 
-        # Run inference
+        # Run model
         with torch.no_grad():
-            mask = model(img_tensor).squeeze().cpu().numpy()
-        
-        # Post-process mask
+            outputs = model(**inputs)
+
+        # Post-process to get mask
+        mask = outputs.logits
+        mask = torch.sigmoid(mask).cpu().numpy()
         mask = (mask > 0.5).astype(np.uint8) * 255
-        mask = Image.fromarray(mask).resize(image.size, Image.LANCZOS)
-        
-        # Apply mask
-        img_rgba = image.convert("RGBA")
-        img_array = np.array(img_rgba)
-        img_array[:, :, 3] = mask
-        result = Image.fromarray(img_array)
-        
-        # Save to bytes buffer
-        buffered = io.BytesIO()
-        result.save(buffered, format="PNG")
-        
-        # Encode as base64
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return f"data:image/png;base64,{img_str}"
-    except Exception as e:
-        return f"Error: {str(e)}"
+        mask = mask.squeeze()
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=remove_background,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Image with Background Removed"),
-    title="DIS Background Removal",
-    description="Remove backgrounds from any image using the open-source DIS (IS-Net) model.",
-    allow_flagging="never"
-)
+        # Apply mask to remove background
+        image_np = np.array(image)
+        alpha_channel = mask
+        result = np.dstack((image_np, alpha_channel))
+        result_image = Image.fromarray(result, mode="RGBA")
+
+        # Save result to bytes
+        output_buffer = io.BytesIO()
+        result_image.save(output_buffer, format="PNG")
+        output_bytes = output_buffer.getvalue()
 
-# Launch the interface
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+        return Response(content=output_bytes, media_type="image/png")
+    except Exception as e:
+        return {"error": str(e)}