Update app.py

app.py

@@ -1,64 +1,86 @@
-import PIL
-import torch
 import gradio as gr
-import os
+import torch
 from process import load_seg_model, get_palette, generate_mask
+from PIL import Image
+import os
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# Initialize model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_path = "model/cloth_segm.pth"
 
-# Initialize models
 try:
-    checkpoint_path = 'model/cloth_segm.pth'
-    if not os.path.exists(checkpoint_path):
-        raise FileNotFoundError(f"Model checkpoint not found at {checkpoint_path}")
-    net = load_seg_model(checkpoint_path, device=device)
+    net = load_seg_model(model_path, device=device)
     palette = get_palette(4)
 except Exception as e:
-    raise RuntimeError(f"Failed to initialize models: {str(e)}")
+    raise RuntimeError(f"Failed to load model: {str(e)}")
 
-def run(img):
-    if img is None:
-        raise gr.Error("Please upload an image first")
+def process_image(input_img):
+    """Process input image and return segmentation mask"""
+    if input_img is None:
+        raise gr.Error("Please upload or capture an image first")
+    
     try:
-        return generate_mask(img, net=net, palette=palette, device=device)
+        # 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
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Handle examples
-image_dir = 'input'
-examples = []
-if os.path.exists(image_dir):
-    examples = [
-        [os.path.join(image_dir, f)] 
-        for f in sorted(os.listdir(image_dir)) 
-        if f.lower().endswith(('.png', '.jpg', '.jpeg'))
-    ]
-
-# Create interface
-with gr.Blocks() as demo:
+# 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.
+    """)
+    
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Input Image", type="pil")
+            input_image = gr.Image(
+                sources=["upload", "webcam"],
+                type="pil",
+                label="Input Image",
+                interactive=True
+            )
+            submit_btn = gr.Button("Process", variant="primary")
+        
         with gr.Column():
-            output_image = gr.Image(label="Segmentation Result")
+            output_image = gr.Image(
+                label="Segmentation Result",
+                interactive=False
+            )
     
-    with gr.Row():
+    # 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=examples,
+            examples=example_images,
             inputs=[input_image],
             outputs=[output_image],
-            fn=run,
+            fn=process_image,
             cache_examples=True,
             label="Example Images"
         )
     
-    submit_btn = gr.Button("Segment", variant="primary")
-    submit_btn.click(fn=run, inputs=input_image, outputs=output_image)
+    submit_btn.click(
+        fn=process_image,
+        inputs=input_image,
+        outputs=output_image
+    )
 
 # Launch with appropriate settings
-try:
-    demo.launch(server_name="0.0.0.0", server_port=7860)
-except Exception as e:
-    print(f"Error launching app: {str(e)}")
-    # Fallback with sharing enabled
-    demo.launch(share=True)
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )