Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| """FinalProject.ipynb | |
| Automatically generated by Colab. | |
| Original file is located at | |
| https://colab.research.google.com/drive/1_wYfP0IRdb9fpc2zvbg8IqdXGx1dTo7X | |
| """ | |
| from datasets import load_dataset | |
| from PIL import Image, ImageChops | |
| from transformers import CLIPProcessor, CLIPModel | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import torch | |
| import numpy as np | |
| import gradio as gr | |
| from diffusers import StableDiffusionImg2ImgPipeline | |
| # Device setup | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Load model and processor | |
| model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device) | |
| processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") | |
| # Load dataset | |
| dataset = load_dataset("lirus18/deepfashion", split="train") | |
| # Precompute image vectors | |
| image_vectors = [] | |
| image_indices = [] | |
| N = 500 # use a smaller subset to avoid long loading | |
| for i in range(N): | |
| img = dataset[i]['image'].convert("RGB") | |
| inputs = processor(images=img, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| emb = model.get_image_features(**inputs) | |
| image_vectors.append(emb.cpu().numpy().squeeze()) | |
| image_indices.append(i) | |
| image_vectors = np.array(image_vectors) | |
| # Find similar images | |
| def find_similar(user_image, top_k=3, exclude_index=None): | |
| inputs = processor(images=user_image.convert("RGB"), return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| query_vec = model.get_image_features(**inputs).cpu().numpy() | |
| sims = cosine_similarity(query_vec, image_vectors)[0] | |
| if exclude_index is not None: | |
| sims[exclude_index] = -1 # Exclude identical | |
| top_idx = sims.argsort()[-top_k:][::-1] | |
| return [dataset[image_indices[i]]['image'] for i in top_idx] | |
| # Load Stable Diffusion pipeline | |
| pipe = StableDiffusionImg2ImgPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| ).to(device) | |
| def generate_outfit_from_image(input_image): | |
| prompt = "fashion outfit design inspired by the clothing item" | |
| init_image = input_image.resize((512, 512)) | |
| result = pipe(prompt=prompt, image=init_image, strength=0.7, guidance_scale=7.5) | |
| return result.images[0] | |
| # Main recommendation function | |
| def recommend_from_upload(uploaded_image): | |
| uploaded_image = uploaded_image.convert("RGB") | |
| closest_idx = None | |
| for i in range(len(image_indices)): | |
| dataset_image = dataset[image_indices[i]]['image'].convert("RGB") | |
| if ImageChops.difference(dataset_image, uploaded_image).getbbox() is None: | |
| closest_idx = i | |
| break | |
| similar_imgs = find_similar(uploaded_image, top_k=3, exclude_index=closest_idx) | |
| generated_img = generate_outfit_from_image(uploaded_image) | |
| return [uploaded_image] + similar_imgs + [generated_img] | |
| # Gradio Interface | |
| demo = gr.Interface( | |
| fn=recommend_from_upload, | |
| inputs=gr.Image(type="pil", label="Upload a clothing item"), | |
| outputs=[ | |
| gr.Image(label="Your Input"), | |
| gr.Image(label="Similar Item 1"), | |
| gr.Image(label="Similar Item 2"), | |
| gr.Image(label="Similar Item 3"), | |
| gr.Image(label="Generated New Outfit"), | |
| ], | |
| title="👗 Fashion Outfit Recommender", | |
| description="Upload a clothing image to see 3 similar outfits and 1 AI-generated one!" | |
| ) | |
| # Only launch if main | |
| if __name__ == "__main__": | |
| demo.launch() | |