Update app.py

app.py

@@ -49,8 +49,11 @@ st.divider()
 
 crop = st.toggle("Activate random crop")
 compress = st.toggle("Activate jpeg compression")
+st.divider()
 
+binary = st.toggle("Activate binary classification")
 st.divider()
+
 file_name = st.file_uploader("Choose an image...")
 
 if file_name is not None:
@@ -70,5 +73,11 @@ if file_name is not None:
     predictions = get_prediction(image)
     
     col2.header("Probabilities")
-    for idx,p in enumerate(predictions):
-        col2.text(f"{ models[idx] }: { round(p * 100, 2)}%")
+    if binary:
+        if predictions[7] > 0.5:
+            col2.markdown(":green[real image!]")
+        else:
+            col2.markdown(":red[fake image!]")
+    else:
+        for idx,p in enumerate(predictions):
+            col2.text(f"{ models[idx] }: { round(p * 100, 2)}%")