Update app.py

app.py

@@ -5,8 +5,9 @@ import torch
 from torchvision.io import read_image
 from transformers import ViTImageProcessor,pipeline
 
-model = ViTImageProcessor.from_pretrained('SeyedAli/Food-Image-Classification-VIT')
-
+# model = ViTImageProcessor.from_pretrained('SeyedAli/Food-Image-Classification-VIT')
+model = ViTForImageClassification.from_pretrained('SeyedAli/Food-Image-Classification-VIT')
+feature_extractor = ViTFeatureExtractor.from_pretrained('SeyedAli/Food-Image-Classification-VIT')
 def FoodClassification(image):
     with tempfile.NamedTemporaryFile(suffix=".png") as temp_image_file:
         # Copy the contents of the uploaded image file to the temporary file
@@ -15,9 +16,13 @@ def FoodClassification(image):
         Image.fromarray(image).save(temp_image_file.name)
         # Load the image file using torchvision
         image = read_image(temp_image_file.name)
-        pipline = pipeline(task="image-classification", model=model)
-        output=pipline(image, return_tensors='pt')
-        return output
+        # Preprocess the image using the ViT feature extractor
+        inputs = feature_extractor(images=image, return_tensors="pt")
+        # Use the ViT model for image classification
+        outputs = model(**inputs)
+        predicted_class_idx = torch.argmax(outputs.logits)
+        predicted_class = model.config.id2label[predicted_class_idx.item()]
+        return predicted_class
 
 iface = gr.Interface(fn=FoodClassification, inputs="image", outputs="text")
 iface.launch(share=False)