Update app.py

app.py

@@ -11,15 +11,12 @@ import cv2
 import chromadb
 from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
 import torch.nn as nn
-import warnings
-
-# Suppress specific warnings
-warnings.filterwarnings("ignore", category=UserWarning, module="transformers.utils.deprecation")
+import matplotlib.pyplot as plt
 
 # Load CLIP model and tokenizer
 @st.cache_resource
 def load_clip_model():
-    model, _, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
+    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
     tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionSigLIP')
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
@@ -30,12 +27,21 @@ clip_model, preprocess_val, tokenizer, device = load_clip_model()
 # Load SegFormer model
 @st.cache_resource
 def load_segformer_model():
-    processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
-    model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+    processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer-b3-fashion")
+    model = AutoModelForSemanticSegmentation.from_pretrained("sayeed99/segformer-b3-fashion")
     return model, processor
 
 segformer_model, segformer_processor = load_segformer_model()
 
+# Load ChromaDB
+@st.cache_resource
+def load_chromadb():
+    client = chromadb.PersistentClient(path="./clothesDB")
+    collection = client.get_collection(name="clothes_items_ver3")
+    return collection
+
+collection = load_chromadb()
+
 # Helper functions
 def load_image_from_url(url, max_retries=3):
     for attempt in range(max_retries):
@@ -50,15 +56,6 @@ def load_image_from_url(url, max_retries=3):
             else:
                 return None
 
-# Load ChromaDB
-@st.cache_resource
-def load_chromadb():
-    client = chromadb.PersistentClient(path="./clothesDB")
-    collection = client.get_collection(name="clothes_items_ver3")
-    return collection
-
-collection = load_chromadb()
-
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
     with torch.no_grad():
@@ -73,35 +70,23 @@ def get_text_embedding(text):
         text_features /= text_features.norm(dim=-1, keepdim=True)
     return text_features.cpu().numpy()
 
-def get_all_embeddings_from_collection(collection):
-    all_embeddings = collection.get(include=['embeddings'])['embeddings']
-    return np.array(all_embeddings)
-
-def get_metadata_from_ids(collection, ids):
-    results = collection.get(ids=ids)
-    return results['metadatas']
-    
 def find_similar_images(query_embedding, collection, top_k=5):
-    database_embeddings = get_all_embeddings_from_collection(collection)
+    database_embeddings = np.array(collection.get(include=['embeddings'])['embeddings'])
     similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
     top_indices = np.argsort(similarities)[::-1][:top_k]
     
     all_data = collection.get(include=['metadatas'])['metadatas']
     
-    top_metadatas = [all_data[idx] for idx in top_indices]
-    
-    results = []
-    for idx, metadata in enumerate(top_metadatas):
-        results.append({
-            'info': metadata,
-            'similarity': similarities[top_indices[idx]]
-        })
+    results = [
+        {'info': all_data[idx], 'similarity': similarities[idx]}
+        for idx in top_indices
+    ]
     return results
 
 def segment_clothing(image):
     inputs = segformer_processor(images=image, return_tensors="pt")
     outputs = segformer_model(**inputs)
-    logits = outputs.logits.cpu()
+    logits = outputs.logits
 
     upsampled_logits = nn.functional.interpolate(
         logits,
@@ -110,7 +95,7 @@ def segment_clothing(image):
         align_corners=False,
     )
 
-    pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
+    pred_seg = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
     
     categories = segformer_model.config.id2label
     segmented_items = []
@@ -127,7 +112,23 @@ def segment_clothing(image):
                     'mask': mask
                 })
     
-    return segmented_items, pred_seg
+    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')
+    
+    # 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]))
@@ -153,7 +154,7 @@ 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 = segment_clothing(query_image)
+                st.session_state.segmentations, st.session_state.pred_seg, st.session_state.categories = segment_clothing(query_image)
                 if st.session_state.segmentations:
                     st.session_state.step = 'select_category'
                 else:
@@ -168,8 +169,10 @@ 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:
-        st.image(st.session_state.pred_seg, caption="Segmentation Map", use_column_width=True)
-    
+        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.subheader("Segmented Clothing Items:")
     
     for segmentation in st.session_state.segmentations: