Update app.py

app.py

@@ -6,26 +6,39 @@ import os
 
 # Title and Description
 st.title("Denim Virtual Try-On")
-st.write("Upload your leg images and the front/back images of denim pants to see how they fit.")
+st.write("Upload your leg image and the front image of denim pants to see how they fit.")
 
 # File Upload Section
 st.sidebar.header("Upload Images")
 subject_image = st.sidebar.file_uploader("Upload Subject Image (Legs)", type=["jpg", "jpeg", "png"])
 pant_front_image = st.sidebar.file_uploader("Upload Denim Pant Front Image", type=["jpg", "jpeg", "png"])
-pant_back_image = st.sidebar.file_uploader("Upload Denim Pant Back Image", type=["jpg", "jpeg", "png"])
+
+# Function to resize images to a manageable size
+def resize_image(image, max_size=(512, 512)):
+    """Resize the image to a max size to reduce payload."""
+    img = Image.open(image)
+    img = img.convert("RGB")  # Ensure image is in RGB format
+    img.thumbnail(max_size)  # Resize while maintaining aspect ratio
+    byte_arr = BytesIO()
+    img.save(byte_arr, format="JPEG")  # Save to BytesIO in JPEG format
+    byte_arr.seek(0)
+    return byte_arr
 
 # Function to process images using Hugging Face API
-def virtual_tryon(subject_img, front_img, back_img):
+def virtual_tryon(subject_img, front_img):
     url = "https://api-inference.huggingface.co/models/ramim36/Kolors-Virtual-Try-On"
     headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}  # Set your Hugging Face API token in env variable
-    
+
+    # Resize images before sending
+    resized_subject = resize_image(subject_img)
+    resized_front = resize_image(front_img)
+
     # Prepare data
     files = {
-        "subject_image": subject_img,
-        "front_image": front_img,
-        "back_image": back_img
+        "subject_image": ("subject.jpg", resized_subject, "image/jpeg"),
+        "front_image": ("front.jpg", resized_front, "image/jpeg")
     }
-    
+
     response = requests.post(url, headers=headers, files=files)
     if response.status_code == 200:
         return Image.open(BytesIO(response.content))
@@ -33,20 +46,30 @@ def virtual_tryon(subject_img, front_img, back_img):
         st.error(f"Error: {response.status_code} - {response.text}")
         return None
 
-# Display Input Images
+# Image Rotation Function
+def rotate_image(image, rotation_angle):
+    """Rotate the uploaded image by the specified angle."""
+    img = Image.open(image)
+    return img.rotate(rotation_angle, expand=True)
+
+# Display Input Images and Add Rotation Option
 if subject_image:
     st.image(subject_image, caption="Uploaded Subject Image", use_column_width=True)
+
 if pant_front_image:
-    st.image(pant_front_image, caption="Uploaded Front Denim Image", use_column_width=True)
-if pant_back_image:
-    st.image(pant_back_image, caption="Uploaded Back Denim Image", use_column_width=True)
+    rotation_angle = st.sidebar.slider("Rotate Front Denim Image", -180, 180, 0)
+    rotated_front_image = rotate_image(pant_front_image, rotation_angle)
+    st.image(rotated_front_image, caption="Rotated Front Denim Image", use_column_width=True)
+    pant_front_image = BytesIO()
+    rotated_front_image.save(pant_front_image, format="JPEG")
+    pant_front_image.seek(0)
 
 # Generate Try-On Output
 if st.button("Generate Virtual Try-On"):
-    if subject_image and pant_front_image and pant_back_image:
+    if subject_image and pant_front_image:
         with st.spinner("Processing... Please wait."):
             try:
-                composite_image = virtual_tryon(subject_image, pant_front_image, pant_back_image)
+                composite_image = virtual_tryon(subject_image, pant_front_image)
                 if composite_image:
                     st.image(composite_image, caption="Virtual Try-On Result", use_column_width=True)
             except Exception as e: