Change model

app.py

@@ -1,46 +1,45 @@
 import gradio as gr
-from transformers import pipeline
+# from transformers import pipeline
 
-# import numpy as np
-# from PIL import Image
+import numpy as np
+from PIL import Image
 
-# from datasets import load_dataset
-# from huggingface_hub import hf_hub_download
-# import tensorflow as tf
-# from tensorflow.keras.applications.resnet50 import preprocess_input
+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 CLASS LABELS FROM DATASET ---
-# ds = load_dataset("dvk65/TrashTypes")
-# class_names = ds["train"].features["label"].names
+# --- LOAD CLASS LABELS FROM DATASET ---
+ds = load_dataset("dvk65/TrashTypes")
+class_names = ds["train"].features["label"].names
 
-classifier = pipeline("image-classification", model="yangy50/garbage-classification")
+# classifier = pipeline("image-classification", model="yangy50/garbage-classification")
 
 
-# # --- LOAD MODEL ---
-# REPO_ID = "dvk65/trash-classifier-resnet50"
-# FILENAME = "trashclassify_13.keras"
+# --- LOAD MODEL ---
+REPO_ID = "dvk65/trash-classifier-resnet50"
+FILENAME = "trashclassify_13.keras"
 
-# model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+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}
-# )
+model = tf.keras.models.load_model(
+    model_path,
+    custom_objects={"preprocess_input": preprocess_input}
+)
 
 
-# def preprocess(image):
-#     image = image.resize((224, 224))
-#     image = np.array(image)
-#     image = preprocess_input(image)
-#     image = np.expand_dims(image, axis=0)
-#     return image
+def preprocess(image):
+    image = image.resize((224, 224))
+    image = np.array(image)
+    image = preprocess_input(image)
+    image = np.expand_dims(image, axis=0)
+    return image
 
 def predict(img):
-    predictions = classifier(img)
-    return {p["label"]: p["score"] for p in predictions}
-    # img = preprocess(img)
-    # preds = model.predict(img)[0]
-    # return {class_names[i]: float(preds[i]) for i in range(len(preds))}
+
+    img = preprocess(img)
+    preds = model.predict(img)[0]
+    return {class_names[i]: float(preds[i]) for i in range(len(preds))}
 
 
 demo = gr.Interface(

requirements.txt

@@ -1,5 +1,6 @@
-
+tensorflow-cpu==2.16.1
+huggingface_hub
+datasets
+gradio
+numpy
 Pillow
-transformers
-torch
-timm