Update app.py

app.py

@@ -21,34 +21,43 @@ class_names = [
 ]
 
 examples = [
-    ['./aeroplane.png'], 
-    ['./horse.png']
+    ['./aeroplane.png'],
+    ['./dog.png'], 
+    ['./horse.png'],
+    ['./ship.png'],
+    ['./truck.png']
 ] 
 
 IMG_SIZE = 72
 
-def infer(input_image):
-    print('#$$$$$$$$$$$$$$$$$$$$$$$$$ IN INFER $$$$$$$$$$$$$$$$$$$$$$$')
-    # image_tensor = read_image(input_image)
-    image_tensor = tf.convert_to_tensor(input_image)
-    image_tensor.set_shape([None, None, 3])
-    image_tensor = tf.image.resize(image_tensor, (IMG_SIZE, IMG_SIZE))
-    print(image_tensor.shape)
+def teacher_model_output(image_tensor):
+    predictions = teacher_model.predict(np.expand_dims((image_tensor), axis=0))
+    predictions = np.squeeze(predictions)
+    predictions = np.argmax(predictions)
+    predicted_label = class_names[predictions.item()]
+    return str(predicted_label)
+    
+def student_model_output(image_tensor):
     predictions = student_model.predict(np.expand_dims((image_tensor), axis=0))
-    print(predictions)
     predictions = np.squeeze(predictions)
-    print(predictions)
-    predictions = np.argmax(predictions) # , axis=2
-    print(predictions)
+    predictions = np.argmax(predictions)
     predicted_label = class_names[predictions.item()]
-    print(predictions.item())
-    print(predicted_label)
     return str(predicted_label)
+
+def infer(input_image):
+    image_tensor = tf.convert_to_tensor(input_image)
+    image_tensor.set_shape([None, None, 3])
+    image_tensor = tf.image.resize(image_tensor, (IMG_SIZE, IMG_SIZE))
+    return teacher_model_output(image_tensor), student_model_output(image_tensor)
     
 input = gr.inputs.Image(shape=(IMG_SIZE, IMG_SIZE))
-output = [gr.outputs.Label()]
-title = "Image Classification using "
-description = "Upload an image or select from examples to classify it.<br>The allowed classes are - Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, Horse, Ship, Truck.<br><p><b>Space author: Vivek Rai</b> <br><b> Keras example author: Sayak Paul </b></p>"
+output = [gr.outputs.Label(label = "Teacher Model Output"), gr.outputs.Label(label = "Student Model Output")]
+
+title = "Image Classification using Consistency training with supervision"
+description = "Upload an image or select from examples to classify it.<br>The allowed classes are - Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, Horse, Ship, Truck.<br><p><b>Teacher Model Repo - </b> <br><b> Student Model Repo -  </b></p>"
+
+
+article = "<div style='text-align: center;'><a href='https://twitter.com/_Blazer_007' target='_blank'>Space by Vivek Rai</a><br><a href='https://keras.io/examples/vision/consistency_training/' target='_blank'>Keras example by Sayak Paul</a></div>"
 
 gr_interface = gr.Interface(
     infer,