Included GAN and Diffusion model also

app.py

@@ -5,12 +5,12 @@ import torch.nn as nn
 import numpy as np
 from torchvision import transforms
 import cv2
-
 from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
 
+# ----------------- Device -----------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ----------------- Load Human Parser Model from Hugging Face Hub -----------------
+# ----------------- Load Human Parser Model -----------------
 processor = AutoImageProcessor.from_pretrained("matei-dorian/segformer-b5-finetuned-human-parsing")
 parser_model = SegformerForSemanticSegmentation.from_pretrained(
     "matei-dorian/segformer-b5-finetuned-human-parsing"
@@ -60,12 +60,6 @@ class UNetGenerator(nn.Module):
         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)),
@@ -85,39 +79,60 @@ 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
+    clothing_labels = [4, 5, 6, 7, 8, 16]
+    for label in clothing_labels:
+        agnostic_np[segmentation_resized == label] = [128, 128, 128]
     return Image.fromarray(agnostic_np)
 
-def generate_tryon_output(agnostic_img, cloth_img):
+def load_model(model_type):
+    model = UNetGenerator().to(device)
+    if model_type == "UNet":
+        checkpoint = torch.load("viton_unet_full_checkpoint.pth", map_location=device)
+    elif model_type == "GAN":
+        checkpoint = torch.load("viton_gan_generator_checkpoint.pth", map_location=device)
+    elif model_type == "Diffusion":
+        checkpoint = torch.load("viton_diffusion_generator_checkpoint.pth", map_location=device)
+    else:
+        raise ValueError("Invalid model type")
+
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model
+
+def generate_tryon_output(agnostic_img, cloth_img, model):
     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 = model(input_tensor)
     output_img = output.squeeze(0).cpu().permute(1, 2, 0).numpy()
+    output_img = (output_img + 1) / 2
     output_img = (output_img * 255).astype(np.uint8)
     return Image.fromarray(output_img)
 
-# ----------------- Gradio Interface -----------------
-def virtual_tryon(person_image, cloth_image):
+# ----------------- Inference Pipeline -----------------
+def virtual_tryon(person_image, cloth_image, model_type):
     segmentation = get_segmentation(person_image)
     agnostic = generate_agnostic(person_image, segmentation)
-    result = generate_tryon_output(agnostic, cloth_image)
+    model = load_model(model_type)
+    result = generate_tryon_output(agnostic, cloth_image, model)
     return agnostic, result
 
+# ----------------- Gradio Interface -----------------
 demo = gr.Interface(
     fn=virtual_tryon,
     inputs=[
         gr.Image(type="pil", label="Person Image"),
-        gr.Image(type="pil", label="Cloth Image")
+        gr.Image(type="pil", label="Cloth Image"),
+        gr.Radio(choices=["UNet", "GAN", "Diffusion"], label="Model Type", value="UNet")
     ],
     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."
+    title="👕 Virtual Try-On App",
+    description="Upload a person image and a clothing image, select a model (UNet, GAN, Diffusion), and try it on virtually."
 )
 
 if __name__ == "__main__":