Update app.py

app.py

@@ -64,7 +64,8 @@ class UNetGenerator(nn.Module):
 # ----------------- Image Transforms -----------------
 img_transform = transforms.Compose([
     transforms.Resize((256, 192)),
-    transforms.ToTensor()
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
 
 # ----------------- Helper Functions -----------------
@@ -108,24 +109,34 @@ def load_model(model_type):
     model.eval()
     return model
 
-def generate_tryon_output(agnostic_img, cloth_img, model):
+def generate_tryon_output(person_img, agnostic_img, cloth_img, segmentation, 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 = model(input_tensor)
-    output_img = output.squeeze(0).cpu().permute(1, 2, 0).numpy()
+
+    output_img = output[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)
+    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)
 
 # ----------------- Inference Pipeline -----------------
 def virtual_tryon(person_image, cloth_image, model_type):
     segmentation = get_segmentation(person_image)
     agnostic = generate_agnostic(person_image, segmentation)
     model = load_model(model_type)
-    result = generate_tryon_output(agnostic, cloth_image, model)
+    result = generate_tryon_output(person_image, agnostic, cloth_image, segmentation, model)
     return agnostic, result
 
 # ----------------- Gradio Interface -----------------