Update app.py

app.py

@@ -9,7 +9,7 @@ import json
 import numpy as np
 import cv2
 import chromadb
-from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+from transformers import pipeline
 import torch.nn as nn
 import matplotlib.pyplot as plt
 
@@ -24,14 +24,12 @@ def load_clip_model():
 
 clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
-# Load SegFormer model
+# Load Clothing Segmentation model
 @st.cache_resource
-def load_segformer_model():
-    processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer-b3-fashion")
-    model = AutoModelForSemanticSegmentation.from_pretrained("sayeed99/segformer-b3-fashion")
-    return model, processor
+def load_segmentation_model():
+    return pipeline(model="mattmdjaga/segformer_b2_clothes")
 
-segformer_model, segformer_processor = load_segformer_model()
+segmenter = load_segmentation_model()
 
 # Load ChromaDB
 @st.cache_resource
@@ -83,55 +81,24 @@ def find_similar_images(query_embedding, collection, top_k=5):
     ]
     return results
 
-def segment_clothing(image):
-    inputs = segformer_processor(images=image, return_tensors="pt")
-    outputs = segformer_model(**inputs)
-    logits = outputs.logits
-
-    upsampled_logits = nn.functional.interpolate(
-        logits,
-        size=image.size[::-1],
-        mode="bilinear",
-        align_corners=False,
-    )
-
-    pred_seg = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
-    
-    categories = segformer_model.config.id2label
-    segmented_items = []
-    
-    for category_id, category_name in categories.items():
-        if category_id in pred_seg:
-            mask = (pred_seg == category_id).astype(np.uint8)
-            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            if contours:
-                x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
-                segmented_items.append({
-                    'category': category_name,
-                    'bbox': [x, y, x+w, y+h],
-                    'mask': mask
-                })
+def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+    segments = segmenter(img)
+    mask_list = []
+    for s in segments:
+        if s['label'] in clothes:
+            mask_list.append(s['mask'])
     
-    return segmented_items, pred_seg, categories
-
-def visualize_segmentation(pred_seg, categories):
-    plt.figure(figsize=(10, 10))
-    plt.imshow(pred_seg, cmap='jet')
-    plt.colorbar(label='Category ID')
-    plt.title("Segmentation Result")
-    plt.axis('off')
+    if mask_list:
+        final_mask = np.array(mask_list[0])
+        for mask in mask_list[1:]:
+            current_mask = np.array(mask)
+            final_mask = final_mask + current_mask
+        
+        final_mask = Image.fromarray(final_mask.astype('uint8') * 255)
+        img = img.convert("RGBA")
+        img.putalpha(final_mask)
     
-    # Add legend
-    unique_classes = np.unique(pred_seg)
-    legend_elements = [plt.Rectangle((0,0),1,1, color=plt.cm.jet(category_id/len(categories)), 
-                       label=f"{category_id}: {categories[category_id]}") 
-                       for category_id in unique_classes if category_id in categories]
-    plt.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1, 0.5))
-    
-    return plt
-
-def crop_image(image, bbox):
-    return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+    return img, segments
 
 # Streamlit app
 st.title("Advanced Fashion Search App")
@@ -154,7 +121,8 @@ if st.session_state.step == 'input':
             query_image = load_image_from_url(st.session_state.query_image_url)
             if query_image is not None:
                 st.session_state.query_image = query_image
-                st.session_state.segmentations, st.session_state.pred_seg, st.session_state.categories = segment_clothing(query_image)
+                segmented_image, st.session_state.segmentations = segment_clothing(query_image)
+                st.session_state.segmented_image = segmented_image
                 if st.session_state.segmentations:
                     st.session_state.step = 'select_category'
                 else:
@@ -169,21 +137,11 @@ elif st.session_state.step == 'select_category':
     with col1:
         st.image(st.session_state.query_image, caption="Original Image", use_column_width=True)
     with col2:
-        seg_fig = visualize_segmentation(st.session_state.pred_seg, st.session_state.categories)
-        st.pyplot(seg_fig)
-        plt.close(seg_fig)  # Prevent memory leaks
+        st.image(st.session_state.segmented_image, caption="Segmented Image", use_column_width=True)
 
     st.subheader("Segmented Clothing Items:")
     
-    for segmentation in st.session_state.segmentations:
-        col1, col2 = st.columns([1, 3])
-        with col1:
-            st.write(f"{segmentation['category']}")
-        with col2:
-            cropped_image = crop_image(st.session_state.query_image, segmentation['bbox'])
-            st.image(cropped_image, caption=segmentation['category'], use_column_width=True)
-    
-    options = [s['category'] for s in st.session_state.segmentations]
+    options = list(set(s['label'] for s in st.session_state.segmentations))
     selected_option = st.selectbox("Select a category to search:", options)
     
     if st.button("Search Similar Items"):
@@ -192,11 +150,15 @@ elif st.session_state.step == 'select_category':
 
 elif st.session_state.step == 'show_results':
     st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
-    selected_segmentation = next(s for s in st.session_state.segmentations 
-                                 if s['category'] == st.session_state.selected_category)
-    cropped_image = crop_image(st.session_state.query_image, selected_segmentation['bbox'])
-    st.image(cropped_image, caption="Cropped Image", use_column_width=True)
-    query_embedding = get_image_embedding(cropped_image)
+    st.image(st.session_state.segmented_image, caption="Segmented Image", use_column_width=True)
+    
+    selected_segment = next(s for s in st.session_state.segmentations if s['label'] == st.session_state.selected_category)
+    mask = np.array(selected_segment['mask'])
+    masked_image = Image.fromarray((np.array(st.session_state.query_image) * mask[:,:,None]).astype('uint8'))
+    
+    st.image(masked_image, caption=f"Selected Category: {st.session_state.selected_category}", use_column_width=True)
+    
+    query_embedding = get_image_embedding(masked_image)
     similar_images = find_similar_images(query_embedding, collection)
     
     st.subheader("Similar Items:")
@@ -220,28 +182,26 @@ elif st.session_state.step == 'show_results':
         st.session_state.segmentations = []
         st.session_state.selected_category = None
 
-else:  # Text search
-    query_text = st.text_input("Enter search text:")
-    if st.button("Search by Text"):
-        if query_text:
-            text_embedding = get_text_embedding(query_text)
-            similar_images = find_similar_images(text_embedding, collection)
-            st.subheader("Similar Items:")
-            for img in similar_images:
-                col1, col2 = st.columns(2)
-                with col1:
-                    st.image(img['info']['image_url'], use_column_width=True)
-                with col2:
-                    st.write(f"Name: {img['info']['name']}")
-                    st.write(f"Brand: {img['info']['brand']}")
-                    category = img['info'].get('category')
-                    if category:
-                        st.write(f"Category: {category}")
-                    st.write(f"Price: {img['info']['price']}")
-                    st.write(f"Discount: {img['info']['discount']}%")
-                    st.write(f"Similarity: {img['similarity']:.2f}")
-        else:
-            st.warning("Please enter a search text.")
+# Text search (optional, you can keep or remove this part)
+st.sidebar.title("Text Search")
+query_text = st.sidebar.text_input("Enter search text:")
+if st.sidebar.button("Search by Text"):
+    if query_text:
+        text_embedding = get_text_embedding(query_text)
+        similar_images = find_similar_images(text_embedding, collection)
+        st.sidebar.subheader("Similar Items:")
+        for img in similar_images:
+            st.sidebar.image(img['info']['image_url'], use_column_width=True)
+            st.sidebar.write(f"Name: {img['info']['name']}")
+            st.sidebar.write(f"Brand: {img['info']['brand']}")
+            category = img['info'].get('category')
+            if category:
+                st.sidebar.write(f"Category: {category}")
+            st.sidebar.write(f"Price: {img['info']['price']}")
+            st.sidebar.write(f"Discount: {img['info']['discount']}%")
+            st.sidebar.write(f"Similarity: {img['similarity']:.2f}")
+    else:
+        st.sidebar.warning("Please enter a search text.")
 
 # Display ChromaDB vacuum message
-st.sidebar.warning("If you've upgraded ChromaDB from a version below 0.6, you may benefit from vacuuming your database. Run 'chromadb utils vacuum --help' for more information.")
+st.sidebar.warning("If you've upgraded ChromaDB from a version below 0.6, you may benefit from vacuuming your database")