Update app.py

app.py

@@ -1,78 +1,88 @@
 import streamlit as st
 from PIL import Image
-import requests
 from io import BytesIO
-import os
+from transformers import pipeline
 
-# Title and Description
-st.title("Denim Virtual Try-On")
-st.write("Upload your leg image and the front image of denim pants to see how they fit.")
+# Title and App Description
+st.title("👗 Virtual Dress Try-On")
+st.write("""
+Upload a **Human Body Image** and a **Garment Image** to generate a Virtual Try-On.  
+Images will be compressed to meet Hugging Face size constraints.
+""")
 
-# 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"])
+# Function to load and compress images
+def compress_image(image_file, max_size_kb=512):
+    """
+    Compresses the input image to meet size constraints while preserving the aspect ratio.
+    Args:
+        image_file: Uploaded file from Streamlit
+        max_size_kb: Maximum file size in kilobytes
+    Returns:
+        Compressed PIL Image object
+    """
+    img = Image.open(image_file).convert("RGB")
+    quality = 95  # Initial compression quality
+    img_format = "JPEG"
+
+    # Compress image iteratively until it meets size constraints
+    while True:
+        img_bytes = BytesIO()
+        img.save(img_bytes, format=img_format, quality=quality)
+        size_kb = len(img_bytes.getvalue()) / 1024  # Size in KB
 
-# 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
+        if size_kb <= max_size_kb or quality <= 10:
+            break
+        quality -= 5  # Reduce quality to compress further
 
-# Function to process images using Hugging Face API
-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')}"}
+    compressed_img = Image.open(img_bytes)
+    return compressed_img
 
-    # Resize images before sending
-    resized_subject = resize_image(subject_img)
-    resized_front = resize_image(front_img)
+# Load Model (Hugging Face Pipeline)
+@st.cache_resource
+def load_model():
+    model_pipeline = pipeline("image-to-image", model="ares1123/virtual-dress-try-on")
+    return model_pipeline
 
-    # Prepare data
-    files = {
-        "subject_image": ("subject.jpg", resized_subject, "image/jpeg"),
-        "front_image": ("front.jpg", resized_front, "image/jpeg")
-    }
+model = load_model()
 
-    response = requests.post(url, headers=headers, files=files)
-    if response.status_code == 200:
-        return Image.open(BytesIO(response.content))
-    else:
-        st.error(f"Error: {response.status_code} - {response.text}")
-        return None
+# Sidebar for Image Upload
+st.sidebar.header("Upload Images")
+uploaded_person = st.sidebar.file_uploader("Upload Human Body Image", type=["jpg", "jpeg", "png"])
+uploaded_clothing = st.sidebar.file_uploader("Upload Garment Image", type=["jpg", "jpeg", "png"])
 
-# 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)
+# Process and Display Images
+if uploaded_person and uploaded_clothing:
+    # Compress uploaded images
+    st.sidebar.info("Compressing images to meet size constraints...")
+    person_image = compress_image(uploaded_person)
+    garment_image = compress_image(uploaded_clothing)
 
-# Display Input Images and Add Rotation Option
-if subject_image:
-    st.image(subject_image, caption="Uploaded Subject Image", use_column_width=True)
+    # Display compressed images
+    col1, col2 = st.columns(2)
+    col1.subheader("Compressed Human Body Image")
+    col1.image(person_image, use_column_width=True)
 
-if pant_front_image:
-    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)
+    col2.subheader("Compressed Garment Image")
+    col2.image(garment_image, use_column_width=True)
 
-# Generate Try-On Output
-if st.button("Generate Virtual Try-On"):
-    if subject_image and pant_front_image:
-        with st.spinner("Processing... Please wait."):
+    # Process button
+    if st.button("👗 Generate Virtual Try-On"):
+        with st.spinner("Processing images... Please wait ⏳"):
             try:
-                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)
+                # Prepare inputs for the model
+                inputs = {"image": person_image, "clothing": garment_image}
+                
+                # Generate output using Hugging Face model
+                output_image = model(inputs)
+                
+                # Display output image
+                st.subheader("✨ Virtual Try-On Result")
+                st.image(output_image, use_column_width=True, caption="Composite Virtual Try-On Image")
             except Exception as e:
-                st.error(f"An error occurred: {e}")
-    else:
-        st.error("Please upload all required images.")
+                st.error(f"An error occurred during processing: {e}")
+else:
+    st.warning("Please upload both the Human Body Image and Garment Image.")
+
+# Footer
+st.markdown("---")
+st.write("Developed with ❤️ using Streamlit and Hugging Face.")