Update app.py

app.py

@@ -1,75 +1,60 @@
-import gradio as gr
+import os
 import torch
-from process import load_seg_model, get_palette, generate_mask
+import gradio as gr
 from PIL import Image
-import os
+from process import load_seg_model, get_palette, generate_mask
 
-# Initialize model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_path = "model/cloth_segm.pth"
+# Device selection
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-try:
-    net = load_seg_model(model_path, device=device)
-    palette = get_palette(4)
-except Exception as e:
-    raise RuntimeError(f"Failed to load model: {str(e)}")
+def load_model():
+    """Load model with Hugging Face Spaces compatible paths"""
+    model_dir = 'model'
+    checkpoint_path = os.path.join(model_dir, 'cloth_segm.pth')
+    
+    # Verify model exists (must be pre-uploaded to HF Spaces)
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(
+            f"Model not found at {checkpoint_path}. "
+            "Please upload the model file to your Space's repository."
+        )
+    
+    try:
+        net = load_seg_model(checkpoint_path, device=device)
+        palette = get_palette(4)
+        return net, palette
+    except Exception as e:
+        raise RuntimeError(f"Model loading failed: {str(e)}")
+
+# Initialize model (will fail fast if there's an issue)
+net, palette = load_model()
 
-def process_image(input_img):
+def process_image(img: Image.Image) -> Image.Image:
     """Process input image and return segmentation mask"""
-    if input_img is None:
-        raise gr.Error("Please upload or capture an image first")
+    if img is None:
+        raise gr.Error("Please upload an image first")
     
     try:
-        # Convert to PIL Image if it's not already
-        if not isinstance(input_img, Image.Image):
-            input_img = Image.fromarray(input_img)
-        
-        # Generate mask
-        output_mask = generate_mask(input_img, net=net, palette=palette, device=device)
-        return output_mask
+        return generate_mask(img, net=net, palette=palette, device=device)
     except Exception as e:
-        raise gr.Error(f"Error processing image: {str(e)}")
+        raise gr.Error(f"Processing failed: {str(e)}")
 
-# Create simple interface
-with gr.Blocks(title="Cloth Segmentation") as demo:
-    gr.Markdown("""
-    # 🧥 Cloth Segmentation App
-    Upload an image or capture from your camera to get segmentation results.
-    """)
+# Gradio interface
+title = "Cloth Segmentation Demo"
+description = """
+Upload an image to get cloth segmentation using U2NET.
+"""
+
+with gr.Blocks() as demo:
+    gr.Markdown(f"## {title}")
+    gr.Markdown(description)
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(
-                sources=["upload", "webcam"],
-                type="pil",
-                label="Input Image",
-                interactive=True
-            )
+            input_image = gr.Image(type="pil", label="Input Image")
             submit_btn = gr.Button("Process", variant="primary")
-        
         with gr.Column():
-            output_image = gr.Image(
-                label="Segmentation Result",
-                interactive=False
-            )
-    
-    # Examples section (optional)
-    example_dir = "input"
-    if os.path.exists(example_dir):
-        example_images = [
-            os.path.join(example_dir, f)
-            for f in os.listdir(example_dir)
-            if f.lower().endswith(('.png', '.jpg', '.jpeg'))
-        ]
-        
-        gr.Examples(
-            examples=example_images,
-            inputs=[input_image],
-            outputs=[output_image],
-            fn=process_image,
-            cache_examples=True,
-            label="Example Images"
-        )
+            output_image = gr.Image(type="pil", label="Segmentation Result")
     
     submit_btn.click(
         fn=process_image,
@@ -77,10 +62,5 @@ with gr.Blocks(title="Cloth Segmentation") as demo:
         outputs=output_image
     )
 
-# Launch with appropriate settings
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    demo.launch()