Update app.py

app.py

@@ -22,6 +22,9 @@ with gr.Blocks() as demo:
     # image classification model for
     # clothes images
     def modelTraining(img):
+        class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
+               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
+        
         # clothing dataset
         mnist = tf.keras.datasets.mnist
 
@@ -72,9 +75,9 @@ with gr.Blocks() as demo:
 
         # Make a prediction using the model
         prediction = probability_model.predict(input_array)
-
+        predicted_label = class_names[np.argmax(prediction)]
         # Postprocess the prediction and return it
-        return np.argmax(prediction)
+        return predicted_label
         
 
     # Creates the Gradio interface objects