make app agnostic to all the models

app.py

@@ -68,17 +68,7 @@ class UNetGenerator(nn.Module):
 #     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 # ])
 #new changes
-if model_type == "UNet":
-    img_transform = transforms.Compose([
-        transforms.Resize((256, 192)),
-        transforms.ToTensor()
-    ])
-else:
-    img_transform = transforms.Compose([
-        transforms.Resize((256, 192)),
-        transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-    ])
+
 #end new changes
 
 # ----------------- Helper Functions -----------------
@@ -144,7 +134,60 @@ def load_model(model_type):
 
 #     return Image.fromarray(final_output)
 #new changes
+# def generate_tryon_output(person_img, agnostic_img, cloth_img, segmentation, model, model_type):
+    
+#     if model_type == "UNet":
+#         img_transform = transforms.Compose([
+#             transforms.Resize((256, 192)),
+#             transforms.ToTensor()
+#         ])
+    
+#     else:
+        
+#         img_transform = transforms.Compose([
+#             transforms.Resize((256, 192)),
+#             transforms.ToTensor(),
+#             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+#         ])
+        
+#     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 = model(input_tensor)
+
+#     if model_type == "UNet":
+#         output_img = output.squeeze(0).cpu().permute(1, 2, 0).numpy()
+#         output_img = (output_img * 255).astype(np.uint8)
+#         return Image.fromarray(output_img)
+#     else:
+#         output_img = output[0].cpu().permute(1, 2, 0).numpy()
+#         output_img = (output_img + 1) / 2
+#         output_img = np.clip(output_img, 0, 1)
+
+#         person_np = np.array(person_img.resize((192, 256))).astype(np.float32) / 255.0
+#         segmentation_resized = cv2.resize(segmentation.astype(np.uint8), (192, 256), interpolation=cv2.INTER_NEAREST)
+#         blend_mask = (segmentation_resized == 0).astype(np.float32)
+#         blend_mask = np.expand_dims(blend_mask, axis=2)
+
+#         final_output = blend_mask * person_np + (1 - blend_mask) * output_img
+#         final_output = (final_output * 255).astype(np.uint8)
+
+#         return Image.fromarray(final_output)
 def generate_tryon_output(person_img, agnostic_img, cloth_img, segmentation, model, model_type):
+    if model_type == "UNet":
+        img_transform = transforms.Compose([
+            transforms.Resize((256, 192)),
+            transforms.ToTensor()
+        ])
+    else:
+        img_transform = transforms.Compose([
+            transforms.Resize((256, 192)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+
     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)
@@ -168,7 +211,6 @@ def generate_tryon_output(person_img, agnostic_img, cloth_img, segmentation, mod
 
         final_output = blend_mask * person_np + (1 - blend_mask) * output_img
         final_output = (final_output * 255).astype(np.uint8)
-
         return Image.fromarray(final_output)
 #new changes end