Update app.py

app.py

@@ -51,18 +51,18 @@ def predict_class(image):
     return prediction[0]
 
 # UI Design for single class output
-def classify_image(image):
-    predicted_class = predict_class(image)
-    output = f"<h2>Predicted Class: <span style='text-transform:uppercase';>{predicted_class}</span></h2>"
-    return output
+# def classify_image(image):
+#     predicted_class = predict_class(image)
+#     output = f"<h2>Predicted Class: <span style='text-transform:uppercase';>{predicted_class}</span></h2>"
+#     return output
 
 
 # UI Design for multy class output
-# def classify_image(image):
-#     results = predict_class(image)
-#     output = {labels.get(i): float(results[i]) for i in range(len(results))}
-#     result = output if max(output.values()) >=0.98 else {"NO_CIFAR10_CLASS": 1}
-#     return result
+def classify_image(image):
+    results = predict_class(image)
+    output = {labels.get(i): float(results[i]) for i in range(len(results))}
+    result = output if max(output.values()) >=0.98 else {"NO_CIFAR10_CLASS": 1}
+    return result[0]
 
 
 inputs = gr.components.Image(type="pil", label="Upload an image")