Change model

app.py

@@ -1,45 +1,40 @@
 import gradio as gr
-import numpy as np
-from PIL import Image
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import torch
 
-from datasets import load_dataset
-from huggingface_hub import hf_hub_download
-import tensorflow as tf
-from tensorflow.keras.applications.resnet50 import preprocess_input
+# Load model + extractor
+model_name = "Aalaa/Fine_tuned_Vit_trash_classification"
 
-# --- LOAD LABELS ---
-ds = load_dataset("dvk65/TrashTypes")
-class_names = ds["train"].features["label"].names
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
 
-# --- LOAD MODEL ---
-REPO_ID = "dvk65/trash-classifier-resnet50"
-FILENAME = "trashclassify_13.keras"
+# Label mapping
+id2label = model.config.id2label
 
-model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+def predict(image):
+    # Convert Gradio PIL to model input
+    inputs = feature_extractor(images=image, return_tensors="pt")
 
-model = tf.keras.models.load_model(
-    model_path,
-    custom_objects={"preprocess_input": preprocess_input}
-)
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+
+    # Return top 3
+    result = {
+        id2label[i]: float(probs[i])
+        for i in probs.topk(3).indices.tolist()
+    }
 
-# --- PREPROCESS---
-def preprocess(image):
-    image = image.resize((224, 224))
-    image = np.array(image).astype("float32")
-    image = np.expand_dims(image, axis=0)
-    return image
+    return result
 
-def predict(img):
-    img = preprocess(img)
-    preds = model.predict(img)[0]
-    
-    return {class_names[i]: float(preds[i]) for i in range(len(preds))}
 
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
-    title="AI Waste Classifier",
+    title="AI Waste Classifier (ViT)"
 )
 
 demo.launch()