commit

utils.py

@@ -10,7 +10,7 @@ import torch.nn.functional as F
 import PIL
 import PIL.Image as Image
 import numpy as np
-
+from transformers import CLIPProcessor, CLIPModel
 
 
 classes_outside_india = ['apple pie', 'baby back ribs', 'baklava', 'beef carpaccio', 'beef tartare',
@@ -40,6 +40,24 @@ classes_india = ['burger','butter_naan', 'chai', 'chapati', 'chole_bhature', 'da
            'pizza', 'samosa']
 
 
+def food_nofood_pred(input_image):
+    # input labels for clip model
+    labels = ['food', 'not food']
+
+        # CLIP Model for classification
+    food_nofood_model = CLIPModel.from_pretrained("flax-community/clip-rsicd-v2")
+    processor = CLIPProcessor.from_pretrained("flax-community/clip-rsicd-v2")
+  
+ #         image = Image.open(requests.get(uploaded_file, stream=True).raw)
+    inputs = processor(text=[f"a photo of a {l}" for l in labels], images=input_image, return_tensors="pt", padding=True)
+    outputs = food_nofood_model(**inputs)
+    logits_per_image = outputs.logits_per_image 
+    probs = logits_per_image.softmax(dim=1)
+    print(probs)
+    pred = probs.detach().cpu().numpy().argmax(axis=1)
+    pred_class = labels[pred[0]]
+    return pred_class
+
 def make_pred_outside_india(input_img, model, device, user_location):
     input_img = input_img.unsqueeze(0)
     model.eval()