Update README.md

README.md

@@ -31,6 +31,7 @@ which use ResNet-50, ResNet-50x4 and ViT32 as image encoders. This repository in
 ```python
 import requests
 from PIL import Image
+import matplotlib.pyplot as plt
 
 from transformers import CLIPProcessor, CLIPModel
 
@@ -39,11 +40,17 @@ processor = CLIPProcessor.from_pretrained("flaviagiammarino/pubmed-clip-vit-base
 
 url = "https://huggingface.co/flaviagiammarino/pubmed-clip-vit-base-patch32/resolve/main/scripts/input.jpeg"
 image = Image.open(requests.get(url, stream=True).raw)
+text = ["Chest X-Ray", "Brain MRI", "Abdominal CT Scan"]
 
-inputs = processor(text=["Chest X-Ray", "Brain MRI", "Abdomen CT Scan"], images=image, return_tensors="pt", padding=True)
-outputs = model(**inputs)
-logits_per_image = outputs.logits_per_image # this is the image-text similarity score
-probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+inputs = processor(text=text, images=image, return_tensors="pt", padding=True)
+probs = model(**inputs).logits_per_image.softmax(dim=1).squeeze()
+
+plt.subplots()
+plt.imshow(image)
+plt.title("".join([x[0] + ": " + x[1] + "\n" for x in zip(text, [format(prob, ".4%") for prob in probs])]))
+plt.axis("off")
+plt.tight_layout()
+plt.show()
 ```
 
 ## Additional Information