Update app.py

app.py

@@ -1,62 +1,43 @@
 import gradio as gr
-import torch
-import torch.nn as nn
-from torchvision import transforms
-from PIL import Image
+import onnxruntime as ort
 import numpy as np
+from PIL import Image
 import requests
-from io import BytesIO
-
-# U-2-Net architecture (simplified, or import from a .py file if you've saved it)
-# You can get the U-2-Net code from https://github.com/xuebinqin/U-2-Net
-
-# For demo, let's download the pre-trained model and use a wrapper instead
-from huggingface_hub import hf_hub_download
-
-# Download u2net.pth from HuggingFace Hub
-model_path = hf_hub_download(repo_id="BritishWerewolf/U-2-Net", filename="onnx/model.onnx")
-
-# Use a known U2NET implementation (e.g., from https://github.com/xuebinqin/U-2-Net/blob/master/u2net_test.py)
-from u2net import U2NET  # Assume you copied the model code as u2net.py
+from torchvision import transforms
 
-# Load model
-model = U2NET(3, 1)
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'), weights_only=False))
-model.eval()
+# Load ONNX model
+ort_session = ort.InferenceSession("model.onnx")  # Ensure model.onnx is in your app folder
 
-# Preprocessing
-transform = transforms.Compose([
-    transforms.Resize((320, 320)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225])
-])
+# Preprocessing function
+def preprocess(image):
+    image = image.resize((320, 320)).convert("RGB")
+    image_np = np.array(image).astype(np.float32) / 255.0
+    image_np = image_np.transpose(2, 0, 1)  # HWC -> CHW
+    image_np = np.expand_dims(image_np, axis=0)  # Add batch dimension
+    return image_np
 
+# Inference + Postprocessing
 def segment_dress(image):
-    original = image.convert("RGB")
-    input_tensor = transform(original).unsqueeze(0)
-    
-    with torch.no_grad():
-        d1, _, _, _, _, _, _ = model(input_tensor)
-        pred = d1[0][0]
-        pred = (pred - pred.min()) / (pred.max() - pred.min())
-        pred_np = pred.cpu().numpy()
+    input_tensor = preprocess(image)
+    inputs = {ort_session.get_inputs()[0].name: input_tensor}
+    outputs = ort_session.run(None, inputs)
     
-    # Resize to original size
-    pred_resized = Image.fromarray((pred_np * 255).astype(np.uint8)).resize(original.size)
-    
-    # Apply mask
-    mask = np.array(pred_resized) / 255.0
-    image_np = np.array(original).astype(np.uint8)
-    segmented = (image_np * mask[..., None]).astype(np.uint8)
+    pred = outputs[0][0][0]
+    pred = (pred - pred.min()) / (pred.max() - pred.min())
+    pred_img = Image.fromarray((pred * 255).astype(np.uint8)).resize(image.size)
+
+    # Apply mask to image
+    image_np = np.array(image.convert("RGB"))
+    mask = np.array(pred_img).astype(np.float32) / 255.0
+    masked = (image_np * mask[..., None]).astype(np.uint8)
 
-    return Image.fromarray(segmented)
+    return Image.fromarray(masked)
 
-# Launch Gradio app
+# Gradio app
 gr.Interface(
     fn=segment_dress,
     inputs=gr.Image(type="pil", label="Upload Image"),
     outputs=gr.Image(type="pil", label="Segmented Dress"),
-    title="Dress Segmentation with U-2-Net",
-    description="Segments the dress (or full foreground) using U-2-Net from Hugging Face"
+    title="U-2-Net Dress Segmentation (ONNX)",
+    description="Upload an image to segment foreground using U-2-Net ONNX model"
 ).launch()