Update app.py

app.py

@@ -9,7 +9,8 @@ import json
 import numpy as np
 import cv2
 import chromadb
-from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+import torch.nn as nn
 
 # Load CLIP model and tokenizer
 @st.cache_resource
@@ -25,8 +26,8 @@ clip_model, preprocess_val, tokenizer, device = load_clip_model()
 # Load SegFormer model
 @st.cache_resource
 def load_segformer_model():
-    model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
     processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+    model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
     return model, processor
 
 segformer_model, segformer_processor = load_segformer_model()
@@ -91,15 +92,23 @@ def find_similar_images(query_embedding, collection, top_k=5):
 def segment_clothing(image):
     inputs = segformer_processor(images=image, return_tensors="pt")
     outputs = segformer_model(**inputs)
-    logits = outputs.logits.squeeze()
-    predicted_mask = logits.argmax(dim=0).numpy()
+    logits = outputs.logits.cpu()
+
+    upsampled_logits = nn.functional.interpolate(
+        logits,
+        size=image.size[::-1],
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
     
     categories = segformer_model.config.id2label
     segmented_items = []
     
     for category_id, category_name in categories.items():
-        if category_id in predicted_mask:
-            mask = (predicted_mask == category_id).astype(np.uint8)
+        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))
@@ -109,7 +118,7 @@ def segment_clothing(image):
                     'mask': mask
                 })
     
-    return segmented_items
+    return segmented_items, pred_seg
 
 def crop_image(image, bbox):
     return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
@@ -135,7 +144,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 = segment_clothing(query_image)
+                st.session_state.segmentations, st.session_state.pred_seg = segment_clothing(query_image)
                 if st.session_state.segmentations:
                     st.session_state.step = 'select_category'
                 else:
@@ -146,7 +155,12 @@ if st.session_state.step == 'input':
             st.warning("Please enter an image URL.")
 
 elif st.session_state.step == 'select_category':
-    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
+    col1, col2 = st.columns(2)
+    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)
+    
     st.subheader("Segmented Clothing Items:")
     
     for segmentation in st.session_state.segmentations: