Update app.py

app.py

@@ -6,7 +6,12 @@ from tensorflow import keras
 
 model = keras.models.load_model('model_CNN.h5')
 
-# Створення функції для передбачення символу
+class_mapping = {1: 'Собака', 2: 'Кінь', 3: 'Слон', 4:'Метелик',
+                 5: 'Курка', 6: 'Кіт', 7:'Корова', 8: 'Вівця',
+                 9: 'Павук', 10: 'Білка'
+                }
+
+  # Створення функції для передбачення тварини
 def predict_image(image):
     image = cv2.resize(image, (224, 224))
     image = np.asarray(image)
@@ -14,7 +19,8 @@ def predict_image(image):
     predictions = model.predict(np.expand_dims(image, axis=0))[0]
     prediction = {}
     for index, probability in enumerate(predictions) :
-      prediction[class_names[index]] = float(round(probability, 3))
+      prediction[class_mapping[index+1]] = float(round(probability, 3))
+      print(prediction)
     return prediction
 
 demo = gr.Blocks()
@@ -22,9 +28,9 @@ demo = gr.Blocks()
 with demo:
     gr.Markdown("Predicte image (dog or cat)")
     with gr.Tab("Predict image"):
-      image_input = gr.Image(label="Upload image")
-      output = gr.Label(label="A letter predicted by a neural network")
-      image_button = gr.Button("Predict")
-      image_button.click(predict_image, inputs=image_input, outputs=output)
+        image_input = gr.Image(label="Upload image")
+        output = gr.Label(label="Animal predicted by neural network")
+        image_button = gr.Button("Predict")
+        image_button.click(predict_image, inputs=image_input, outputs=output)
 
-demo.launch()
+demo.launch()