Update app.py

app.py

@@ -6,6 +6,26 @@ import numpy as np
 from PIL import Image
 import tensorflow as tf
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+import requests
+from io import BytesIO
+
+model_name = "facebook/deit-base-distilled-patch16"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+
+def classify_image(image):
+    # Load and preprocess the image
+    image = Image.open(BytesIO(image))
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    
+    # Perform image classification
+    with torch.no_grad():
+        outputs = model(**inputs)
+    predicted_class = outputs.logits.argmax().item()
+    
+    return model.config.id2label[predicted_class]
+
 
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
     "nvidia/segformer-b0-finetuned-cityscapes-512-1024"
@@ -102,7 +122,8 @@ def sepia(input_img):
     return fig
 
 
-demo = gr.Interface(fn=sepia,
+
+demo = gr.Interface(fn=classify_image,
                     inputs=gr.Image(shape=(400, 600)),
                     outputs=['plot'],
                     title="SWJIN11 TASK",