feat: add cosine similarity measure

app.py

@@ -3,6 +3,7 @@ from lib.utils.model import get_model, get_similarities
 from PIL import Image 
 
 st.title('IRRA Text-To-Image-Retrival')
+st.markdown('A text-to-image retrieval model implemented from [arXiv: Cross-Modal Implicit Relation Reasoning and Aligning for Text-to-Image Person Retrieval](https://arxiv.org/abs/2303.12501)')
 
 st.header('Inputs')
 caption = st.text_input('Description Input')
@@ -12,7 +13,7 @@ if images is not None:
     st.image(images) # type: ignore
 
 st.header('Options')
-st.subheader('Ranks')
+st.subheader('Ranks', help='How many predictions the model is allowed to make')
 
 ranks = st.slider('slider_ranks', min_value=1, max_value=10, label_visibility='collapsed',value=5)
 
@@ -26,15 +27,16 @@ if button:
     st.text(f'IRRA model loaded with {sum(p.numel() for p in model.parameters()) / 1e6:.0f}M parameters')
     
     with st.spinner('Computing and ranking similarities'):
-        similarities = get_similarities(caption, images, model)
+        similarities = get_similarities(caption, images, model).squeeze(0)
 
-    indices = similarities.argsort(descending=True).squeeze(0).cpu().tolist()[:ranks]
+    indices = similarities.argsort(descending=True).cpu().tolist()[:ranks]
     
     for i, idx in enumerate(indices):
-        c1, c2 = st.columns(2)
+        c1, c2, c3 = st.columns(3)
         with c1:
             st.text(f'Rank {i + 1}')
         with c2:
             st.image(images[idx])
-
+        with c3:
+            st.text(f'Cosine sim {similarities[idx].cpu():.2f}')
     

lib/utils/model.py

@@ -24,7 +24,6 @@ def get_similarities(text: str, images: list[str], model: IRRA) -> torch.Tensor:
     txt = tokenize(text, tokenizer)
     imgs = prepare_images(images)
 
-    print(imgs.shape)
     image_feats = model.encode_image(imgs)
     text_feats = model.encode_text(txt.unsqueeze(0))