paloma99 commited on
Commit
1a8d06e
1 Parent(s): 5b3d8e7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -4
app.py CHANGED
@@ -40,12 +40,41 @@ import shutil
40
  custom_title = "<span style='color: rgb(243, 239, 224);'>Green Greta</span>"
41
 
42
 
 
 
 
 
 
 
43
  # Cell 1: Image Classification Model
44
- image_pipeline = pipeline(task="image-classification", model="rocioadlc/EfficientNetV2L")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
45
 
46
- def predict_image(input_img):
47
- predictions = image_pipeline(input_img)
48
- return {p["label"]: p["score"] for p in predictions}
49
 
50
  image_gradio_app = gr.Interface(
51
  fn=predict_image,
 
40
  custom_title = "<span style='color: rgb(243, 239, 224);'>Green Greta</span>"
41
 
42
 
43
+ from huggingface_hub import from_pretrained_keras
44
+
45
+ import tensorflow as tf
46
+ from tensorflow import keras
47
+ from PIL import Image
48
+
49
  # Cell 1: Image Classification Model
50
+ model1 = from_pretrained_keras("rocioadlc/EfficientNetV2L")
51
+
52
+ # Define class labels
53
+ class_labels = ['cardboard', 'compost', 'glass', 'metal', 'paper', 'plastic', 'trash']
54
+
55
+ # Function to predict image label and score
56
+ def predict_image(input):
57
+ # Resize the image to the size expected by the model
58
+ image = input.resize((224, 224))
59
+ # Convert the image to a NumPy array
60
+ image_array = tf.keras.preprocessing.image.img_to_array(image)
61
+ # Normalize the image
62
+ image_array /= 255.0
63
+ # Expand the dimensions to create a batch
64
+ image_array = tf.expand_dims(image_array, 0)
65
+ # Predict using the model
66
+ predictions = model1.predict(image_array)
67
+
68
+ # Get the predicted class label
69
+ predicted_class_index = tf.argmax(predictions, axis=1).numpy()[0]
70
+ predicted_class_label = class_labels[predicted_class_index]
71
+
72
+ # Get the confidence score of the predicted class
73
+ confidence_score = predictions[0][predicted_class_index]
74
+
75
+ # Return input image path, predicted class label, and confidence score
76
+ return input, {predicted_class_label: confidence_score}
77
 
 
 
 
78
 
79
  image_gradio_app = gr.Interface(
80
  fn=predict_image,