Update app.py

app.py

@@ -4,55 +4,79 @@ from process import load_seg_model, get_palette, generate_mask
 from PIL import Image
 import os
 
-# Model initialization
+# Initialize model
 device = "cuda" if torch.cuda.is_available() else "cpu"
+model_path = "model/cloth_segm.pth"
 
-def load_models():
-    try:
-        net = load_seg_model("model/cloth_segm.pth", device=device)
-        palette = get_palette(4)
-        return net, palette
-    except Exception as e:
-        raise gr.Error(f"Model failed to load: {str(e)}")
+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)}")
 
-net, palette = load_models()
-
-def predict(image):
-    if image is None:
+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:
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        return generate_mask(image, 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"Processing error: {str(e)}")
+        raise gr.Error(f"Error processing image: {str(e)}")
 
-# Interface
+# Create simple interface
 with gr.Blocks(title="Cloth Segmentation") as demo:
-    gr.Markdown("## 👕 Cloth Segmentation Tool")
+    gr.Markdown("""
+    # 🧥 Cloth Segmentation App
+    Upload an image or capture from your camera to get segmentation results.
+    """)
     
     with gr.Row():
         with gr.Column():
-            img_input = gr.Image(sources=["upload", "webcam"], 
-                               type="pil", 
-                               label="Input Image")
-            btn = gr.Button("Generate Mask", variant="primary")
+            input_image = gr.Image(
+                sources=["upload", "webcam"],
+                type="pil",
+                label="Input Image",
+                interactive=True
+            )
+            submit_btn = gr.Button("Process", variant="primary")
         
         with gr.Column():
-            img_output = gr.Image(label="Segmentation Result")
+            output_image = gr.Image(
+                label="Segmentation Result",
+                interactive=False
+            )
     
-    # Optional examples
-    if os.path.exists("examples"):
+    # 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=[os.path.join("examples", f) for f in os.listdir("examples") 
-                     if f.endswith(('.png','.jpg','.jpeg'))],
-            inputs=img_input,
-            outputs=img_output,
-            fn=predict,
-            cache_examples=True
+            examples=example_images,
+            inputs=[input_image],
+            outputs=[output_image],
+            fn=process_image,
+            cache_examples=True,
+            label="Example Images"
         )
+    
+    submit_btn.click(
+        fn=process_image,
+        inputs=input_image,
+        outputs=output_image
+    )
 
-    btn.click(predict, inputs=img_input, outputs=img_output)
-
+# Launch with appropriate settings
 if __name__ == "__main__":
     demo.launch()