Change

app.py

@@ -1,29 +1,28 @@
 import gradio as gr
-# from transformers import pipeline
-
 import numpy as np
 from PIL import Image
 
 from datasets import load_dataset
 from huggingface_hub import hf_hub_download
 import tensorflow as tf
-import tensorflow as tf
+from tensorflow.keras.applications.resnet50 import preprocess_input
 
-# --- LOAD CLASS LABELS FROM DATASET ---
+# --- LOAD LABELS ---
 ds = load_dataset("dvk65/TrashTypes")
 class_names = ds["train"].features["label"].names
 
-# classifier = pipeline("image-classification", model="yangy50/garbage-classification")
-
-
 # --- LOAD MODEL ---
 REPO_ID = "dvk65/trash-classifier-resnet50"
 FILENAME = "trashclassify_13.keras"
-model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
 
-model = tf.keras.models.load_model(model_path)
+model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
 
+model = tf.keras.models.load_model(
+    model_path,
+    custom_objects={"preprocess_input": preprocess_input}
+)
 
+# --- PREPROCESS---
 def preprocess(image):
     image = image.resize((224, 224))
     image = np.array(image).astype("float32")
@@ -35,7 +34,6 @@ def predict(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"),