Upload app.py with huggingface_hub

app.py

@@ -1,121 +1,69 @@
 """
 HuggingFace Space for Clothing Segmentation
-Uses DeepLabV3+ ONNX model to segment human body parts
 """
 
 import os
+from pathlib import Path
 import numpy as np
 from PIL import Image
-import onnxruntime
 import gradio as gr
-from gradio import processing_utils
+from huggingface_hub import hf_hub_download
 
-# Labels for the model
-LABELS = [
-    "background", "unknown", "hair", "unknown", "glasses",
-    "top-clothes", "unknown", "unknown", "unknown", "bottom-clothes",
-    "torso-skin", "unknown", "unknown", "face", "left-arm",
-    "right-arm", "left-leg", "right-leg", "left-foot", "right-foot"
-]
-
-# Clothing-related classes (top-clothes, bottom-clothes)
+# Clothing classes from the model
 CLOTHING_CLASSES = [5, 9]  # top-clothes, bottom-clothes
 
-def load_model():
-    """Load ONNX model"""
-    model_path = "deeplabv3p-resnet50-human.onnx"
-    if not os.path.exists(model_path):
-        # Try to download from HuggingFace
-        from huggingface_hub import hf_hub_download
-        model_path = hf_hub_download(
-            repo_id="Metal3d/deeplabv3p-resnet50-human",
-            filename="deeplabv3p-resnet50-human.onnx"
-        )
-    session = onnxruntime.InferenceSession(model_path)
-    return session
+print("Downloading model...")
+model_path = hf_hub_download(
+    repo_id="Metal3d/deeplabv3p-resnet50-human",
+    filename="deeplabv3p-resnet50-human.onnx"
+)
+print(f"Model downloaded to: {model_path}")
 
-# Load model at startup
-print("Loading model...")
-model = load_model()
+import onnxruntime as ort
+session = ort.InferenceSession(model_path)
 print("Model loaded!")
 
-def preprocess_image(img: Image.Image) -> np.ndarray:
+def preprocess(img):
     """Preprocess image for model"""
     img = img.resize((512, 512))
-    img_array = np.array(img).astype(np.float32) / 127.5 - 1
-    # Ensure 3 channels
-    if len(img_array.shape) == 2:
-        img_array = np.stack([img_array] * 3, axis=-1)
-    elif img_array.shape[-1] == 4:
-        img_array = img_array[:, :, :3]
-    return img_array
+    arr = np.array(img).astype(np.float32) / 127.5 - 1
+    if len(arr.shape) == 2:
+        arr = np.stack([arr] * 3, axis=-1)
+    elif arr.shape[-1] == 4:
+        arr = arr[:, :, :3]
+    return np.transpose(arr, (2, 0, 1))[np.newaxis, :, :, :]
 
-def segment_clothing(img: Image.Image) -> Image.Image:
-    """Segment clothing from image"""
-    # Preprocess
-    input_data = preprocess_image(img)
-    input_data = np.transpose(input_data, (2, 0, 1))[np.newaxis, :, :, :]
-    
-    # Run inference
-    input_name = model.get_inputs()[0].name
-    output_name = model.get_outputs()[0].name
-    result = model.run([output_name], {input_name: input_data[0]})[0]
+def process(user_img, fabric_img):
+    """Process images"""
+    if user_img is None or fabric_img is None:
+        return None
     
-    # Get argmax
+    input_data = preprocess(user_img)
+    input_name = session.get_inputs()[0].name
+    output_name = session.get_outputs()[0].name
+    result = session.run([output_name], {input_name: input_data[0]})[0]
     result = np.argmax(result[0], axis=0)
+    mask = np.isin(result, CLOTHING_CLASSES).astype(np.uint8) * 255
+    mask_img = Image.fromarray(mask).resize(user_img.size, Image.NEAREST)
     
-    # Create clothing mask (top + bottom clothes)
-    clothing_mask = np.isin(result, CLOTHING_CLASSES).astype(np.uint8) * 255
-    
-    # Resize back to original
-    mask_img = Image.fromarray(clothing_mask).resize(img.size, Image.NEAREST)
-    return mask_img
-
-def apply_fabric(user_img: Image.Image, fabric_img: Image.Image) -> Image.Image:
-    """Apply fabric to segmented clothing area"""
-    # First segment the clothing
-    mask = segment_clothing(user_img)
-    
-    # Convert to numpy
-    user_arr = np.array(user_img)
     fabric_arr = np.array(fabric_img.resize(user_img.size, Image.LANCZOS))
-    mask_arr = np.array(mask) / 255.0
+    user_arr = np.array(user_img)
+    mask_arr = np.array(mask_img) / 255.0
     
-    # Apply: where mask=1, use fabric; where mask=0, use original
-    result = (fabric_arr * mask_arr[:, :, np.newaxis] + 
-              user_arr * (1 - mask_arr[:, :, np.newaxis])).astype(np.uint8)
+    output = (fabric_arr * mask_arr[:, :, np.newaxis] + 
+             user_arr * (1 - mask_arr[:, :, np.newaxis])).astype(np.uint8)
     
-    return Image.fromarray(result)
+    return Image.fromarray(output)
 
-# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Clothing Virtual Try-On")
-    gr.Markdown("Upload your photo and select a fabric to try it on!")
-    
+    gr.Markdown("# 👗 Virtual Try-On")
     with gr.Row():
         with gr.Column():
-            user_image = gr.Image(type="pil", label="Your Photo")
-            fabric_image = gr.Image(type="pil", label="Fabric")
-            
+            user = gr.Image(type="pil", label="Your Photo")
+            fabric = gr.Image(type="pil", label="Fabric")
         with gr.Column():
-            output_image = gr.Image(type="pil", label="Result")
-    
-    with gr.Row():
-        submit_btn = gr.Button("Apply Fabric", variant="primary")
-    
-    submit_btn.click(
-        fn=apply_fabric,
-        inputs=[user_image, fabric_image],
-        outputs=output_image
-    )
+            result = gr.Image(type="pil", label="Result")
     
-    gr.Examples(
-        examples=[
-            ["https://example.com/person.jpg", "https://example.com/fabric.jpg"],
-        ],
-        inputs=[user_image, fabric_image],
-        outputs=output_image,
-    )
+    gr.Button("Apply").click(fn=process, inputs=[user, fabric], outputs=result)
 
-# Launch space
 demo.launch(server_port=7860)