Update app.py

app.py

@@ -5,36 +5,78 @@ import requests
 import numpy as np
 from PIL import Image
 
-shapes = ["leggings", "jogger", 
-          "palazzo", "cargo",
-          "dresspants", "chinos"]
+catgs = [
+  "Shirts",
+  "SetShirtsPants",
+  "SetJacketsPants",
+  "Pants",
+  "Jeans",
+  "JacketsCoats",
+  "Shoes",
+  "Underpants",
+  "Socks",
+  "Hats",
+  "Wallets",
+  "Bags",
+  "Scarfs",
+  "Parasols&Umbrellas",
+  "Necklaces",
+  "Towels&Robes",
+  "WallObjects",
+  "Rugs",
+  "Glassware",
+  "Mugs&Cups",
+  "OralCare"
+]
 
-model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionCLIP')
-tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionCLIP')
+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')
 
-
-shapes_desc = list(map(lambda x: "a " + x + " pants shape", shapes))
-text = tokenizer(shapes_desc)
+text = tokenizer(catgs)
 
 with torch.no_grad(), torch.cuda.amp.autocast():
-  text_features = model.encode_text(text)
-  text_features /= text_features.norm(dim=-1, keepdim=True)
+    text_features = model.encode_text(text)
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+
 
 def predict(inp):
   image = preprocess_val(inp).unsqueeze(0)
 
-  with torch.no_grad(), torch.cuda.amp.autocast():
-      image_features = model.encode_image(image)
-      image_features /= image_features.norm(dim=-1, keepdim=True)
-
-      text_probs = (100 * image_features @ text_features.T).softmax(dim=-1)
-      confidences = {shapes[i]: float(text_probs[0, i]) for i in range(6)}
-  return confidences
-
+    with torch.no_grad(), torch.cuda.amp.autocast():
+        image_features = model.encode_image(image)
+        image_features /= image_features.norm(dim=-1, keepdim=True)    
+        text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+    
+    max_prob_idx = np.argmax(text_probs)
+    pred_lbl = catgs[max_prob_idx]
+    pred_lbl_prob = text_probs[0, max_prob_idx].item()
+    
+    mw = ["men", "women", "boy", "girl"]
+    catgs = [
+        mw[0] + "s " + pred_lbl,
+        mw[1] + "s " + pred_lbl,
+        mw[2] + "s " + pred_lbl,
+        mw[3] + "s " + pred_lbl
+    ]
+    text = tokenizer(catgs)
+    
+    with torch.no_grad(), torch.cuda.amp.autocast():
+        image_features = model.encode_image(image)
+        text_features = model.encode_text(text)
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+    
+        text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+    
+    max_prob_idx = np.argmax(text_probs)
+    pred_lbl_f = mw[max_prob_idx]
+    pred_lbl_prob_f = text_probs[0, max_prob_idx].item()
+    tlt = f"{pred_lbl} <{100.0 * pred_lbl_prob:.1f}%> , {pred_lbl_f} <{100.0 * pred_lbl_prob_f:.1f}%>"
+    return(tlt)
 
 gr.Interface(fn=predict,
              inputs=gr.Image(type="pil"),
-             outputs=gr.Label(num_top_classes=6),
+             outputs=gr.Label(),
              examples=["imgs/cargo.jpg", "imgs/palazzo.jpg",
                       "imgs/leggings.jpg", "imgs/jogger.jpg",
                       "imgs/chinos.jpg", "imgs/dresspants.jpg"]).launch(share=True)