Create app.py

app.py

@@ -0,0 +1,80 @@
+import streamlit as st
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+from torchvision import models, transforms
+
+# Load a pre-trained body segmentation model (DeepLabV3) for human segmentation
+model = models.segmentation.deeplabv3_resnet101(pretrained=True)
+model.eval()
+
+# Define the transformation for the input image
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Function to segment the body
+def segment_body(image):
+    # Convert the image to RGB (needed by the model)
+    image = image.convert("RGB")
+    
+    # Apply the transformations
+    input_tensor = transform(image).unsqueeze(0)
+    
+    # Get the model's prediction
+    with torch.no_grad():
+        output = model(input_tensor)["out"][0]
+    
+    # Convert the output to a probability map and get the human mask (class 15 is human)
+    output_predictions = output.argmax(0)
+    human_mask = output_predictions == 15  # Class 15 corresponds to human
+    
+    return human_mask
+
+# Function to overlay clothing onto the image
+def overlay_clothing(user_image, clothing_image, mask):
+    # Resize clothing image to match the user's body size (for simplicity, we use resizing here)
+    clothing_resized = clothing_image.resize((user_image.width, user_image.height), Image.ANTIALIAS)
+    
+    # Convert to numpy arrays for OpenCV processing
+    user_image_np = np.array(user_image)
+    clothing_resized_np = np.array(clothing_resized)
+    
+    # Apply the mask to the clothing image and the user's body image
+    clothing_overlay = np.zeros_like(user_image_np)
+    clothing_overlay[mask] = clothing_resized_np[mask]
+    
+    # Combine the user's image with the clothing overlay
+    result = user_image_np.copy()
+    result[mask] = clothing_overlay[mask]
+    
+    return Image.fromarray(result)
+
+def main():
+    st.title("Virtual Try-On App")
+    st.write("Upload a photo of yourself and a clothing item to virtually try it on!")
+    
+    # Upload user image
+    uploaded_user_image = st.file_uploader("Upload your photo", type=["jpg", "png", "jpeg"])
+    
+    # Upload clothing item
+    uploaded_clothing_image = st.file_uploader("Upload the clothing item (e.g., T-shirt)", type=["jpg", "png", "jpeg"])
+
+    if uploaded_user_image and uploaded_clothing_image:
+        # Open images
+        user_image = Image.open(uploaded_user_image)
+        clothing_image = Image.open(uploaded_clothing_image)
+
+        # Segment the user's body from the image
+        mask = segment_body(user_image)
+
+        # Overlay the clothing item onto the user's body
+        result_image = overlay_clothing(user_image, clothing_image, mask)
+
+        # Show the result image
+        st.image(result_image, caption="Virtual Try-On Result", use_column_width=True)
+
+if __name__ == "__main__":
+    main()