Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PIL import Image | |
| from io import BytesIO | |
| from transformers import pipeline | |
| # 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. | |
| """) | |
| # 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 | |
| if size_kb <= max_size_kb or quality <= 10: | |
| break | |
| quality -= 5 # Reduce quality to compress further | |
| compressed_img = Image.open(img_bytes) | |
| return compressed_img | |
| # Load Model (Hugging Face Pipeline) | |
| def load_model(): | |
| model_pipeline = pipeline("image-to-image", model="ares1123/virtual-dress-try-on") | |
| return model_pipeline | |
| model = load_model() | |
| # 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"]) | |
| # 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 compressed images | |
| col1, col2 = st.columns(2) | |
| col1.subheader("Compressed Human Body Image") | |
| col1.image(person_image, use_column_width=True) | |
| col2.subheader("Compressed Garment Image") | |
| col2.image(garment_image, use_column_width=True) | |
| # Process button | |
| if st.button("👗 Generate Virtual Try-On"): | |
| with st.spinner("Processing images... Please wait ⏳"): | |
| try: | |
| # 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 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.") | |