Create app.py

app.py

@@ -0,0 +1,124 @@
+import gradio as gr
+from PIL import Image
+import torch
+import torch.nn as nn
+import numpy as np
+from torchvision import transforms
+import cv2
+
+from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ----------------- Load Human Parser Model from Hugging Face Hub -----------------
+processor = AutoImageProcessor.from_pretrained("matei-dorian/segformer-b5-finetuned-human-parsing")
+parser_model = SegformerForSemanticSegmentation.from_pretrained(
+    "matei-dorian/segformer-b5-finetuned-human-parsing"
+).to(device).eval()
+
+# ----------------- UNet Generator Definition -----------------
+class UNetGenerator(nn.Module):
+    def __init__(self, in_channels=6, out_channels=3):
+        super(UNetGenerator, self).__init__()
+
+        def block(in_c, out_c):
+            return nn.Sequential(
+                nn.Conv2d(in_c, out_c, 4, 2, 1),
+                nn.BatchNorm2d(out_c),
+                nn.ReLU(inplace=True)
+            )
+
+        def up_block(in_c, out_c):
+            return nn.Sequential(
+                nn.ConvTranspose2d(in_c, out_c, 4, 2, 1),
+                nn.BatchNorm2d(out_c),
+                nn.ReLU(inplace=True)
+            )
+
+        self.down1 = block(in_channels, 64)
+        self.down2 = block(64, 128)
+        self.down3 = block(128, 256)
+        self.down4 = block(256, 512)
+
+        self.up1 = up_block(512, 256)
+        self.up2 = up_block(512, 128)
+        self.up3 = up_block(256, 64)
+        self.up4 = nn.Sequential(
+            nn.ConvTranspose2d(128, out_channels, 4, 2, 1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        d1 = self.down1(x)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+
+        u1 = self.up1(d4)
+        u2 = self.up2(torch.cat([u1, d3], dim=1))
+        u3 = self.up3(torch.cat([u2, d2], dim=1))
+        u4 = self.up4(torch.cat([u3, d1], dim=1))
+        return u4
+
+# ----------------- Load UNet Try-On Model -----------------
+tryon_model = UNetGenerator().to(device)
+checkpoint = torch.load("viton_unet_full_checkpoint.pth", map_location=device)
+tryon_model.load_state_dict(checkpoint['model_state_dict'])
+tryon_model.eval()
+
+# ----------------- Image Transforms -----------------
+img_transform = transforms.Compose([
+    transforms.Resize((256, 192)),
+    transforms.ToTensor()
+])
+
+# ----------------- Helper Functions -----------------
+def get_segmentation(image: Image.Image):
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = parser_model(**inputs)
+    logits = outputs.logits
+    predicted = torch.argmax(logits, dim=1)[0].cpu().numpy()
+    return predicted
+
+def generate_agnostic(image: Image.Image, segmentation):
+    img_np = np.array(image.resize((192, 256)))
+    agnostic_np = img_np.copy()
+    segmentation_resized = cv2.resize(segmentation.astype(np.uint8), (192, 256), interpolation=cv2.INTER_NEAREST)
+    agnostic_np[segmentation_resized == 4] = [128, 128, 128]  # Mask upper clothes
+    return Image.fromarray(agnostic_np)
+
+def generate_tryon_output(agnostic_img, cloth_img):
+    agnostic_tensor = img_transform(agnostic_img).unsqueeze(0).to(device)
+    cloth_tensor = img_transform(cloth_img).unsqueeze(0).to(device)
+    input_tensor = torch.cat([agnostic_tensor, cloth_tensor], dim=1)
+
+    with torch.no_grad():
+        output = tryon_model(input_tensor)
+    output_img = output.squeeze(0).cpu().permute(1, 2, 0).numpy()
+    output_img = (output_img * 255).astype(np.uint8)
+    return Image.fromarray(output_img)
+
+# ----------------- Gradio Interface -----------------
+def virtual_tryon(person_image, cloth_image):
+    segmentation = get_segmentation(person_image)
+    agnostic = generate_agnostic(person_image, segmentation)
+    result = generate_tryon_output(agnostic, cloth_image)
+    return agnostic, result
+
+demo = gr.Interface(
+    fn=virtual_tryon,
+    inputs=[
+        gr.Image(type="pil", label="Person Image"),
+        gr.Image(type="pil", label="Cloth Image")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Agnostic (Torso Masked)"),
+        gr.Image(type="pil", label="Virtual Try-On Output")
+    ],
+    title="👕 Virtual Try-On (UNet + Segformer)",
+    description="Upload a person image and a cloth image to try on the cloth virtually."
+)
+
+if __name__ == "__main__":
+    demo.launch()