Update app.py

app.py

@@ -40,12 +40,49 @@ import shutil
 custom_title = "<span style='color: rgb(243, 239, 224);'>Green Greta</span>"
 
 
-# Cell 1: Image Classification Model
-image_pipeline = pipeline(task="image-classification", model="guillen/vit-basura-test1")
+from huggingface_hub import from_pretrained_keras
+
+import tensorflow as tf
+from tensorflow import keras
+from PIL import Image
 
-def predict_image(input_img):
-    predictions = image_pipeline(input_img)
-    return {p["label"]: p["score"] for p in predictions} 
+# Cell 1: Image Classification Model
+model1 = from_pretrained_keras("rocioadlc/EfficientNetV2L")
+
+class_labels = ['battery',
+ 'biological',
+ 'brown-glass',
+ 'cardboard',
+ 'clothes',
+ 'green-glass',
+ 'metal',
+ 'paper',
+ 'plastic',
+ 'shoes',
+ 'trash',
+ 'white-glass']
+
+def predict_image(input):
+    # Resize the image to the size expected by the model
+    image = input.resize((244, 224))
+    # Convert the image to a NumPy array
+    image_array = tf.keras.preprocessing.image.img_to_array(image)
+    # Normalize the image
+    image_array /= 255.0
+    # Expand the dimensions to create a batch
+    image_array = tf.expand_dims(image_array, 0)
+    # Predict using the model
+    predictions = model1.predict(image_array)
+    
+    # Get the predicted class label
+    predicted_class_index = np.argmax(predictions, -1)
+    predicted_class_label = class_labels[predicted_class_index]
+    
+    # Get the confidence score of the predicted class
+    confidence_score = predictions[0][predicted_class_index]
+    
+    # Return predicted class label and confidence score
+    return {predicted_class_label: confidence_score}
 
 image_gradio_app = gr.Interface(
     fn=predict_image,