Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+# CHANGED: We now import AutoImageProcessor instead of AutoFeatureExtractor
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import torch
 
@@ -7,14 +8,16 @@ import torch
 MODEL_NAME = "prithivMLmods/Augmented-Waste-Classifier-SigLIP2"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
+# CHANGED: Use AutoImageProcessor
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(device)
 
 # Inference function
 def classify_image(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
-    inputs = extractor(images=image, return_tensors="pt")
+    # CHANGED: Use the new 'processor' variable
+    inputs = processor(images=image, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
     outputs = model(**inputs)
     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
@@ -37,4 +40,4 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()