Update app.py

app.py

@@ -4,18 +4,21 @@ model_0 = tf.keras.models.load_model('bestmodel.h5')
 def classify_image(inp):
     inp = inp.reshape((-1, 224, 224, 3))
     prediction = model_0.predict(inp)
-    if prediction.argmax() == 0:
-      output = "Rifle violence"
-    elif prediction.argmax() == 1:
-      output = "guns violence"
-    elif prediction.argmax() == 2:
-      output = "knife violence"
-    elif prediction.argmax() == 3:
-      output = "image porno"
-    elif prediction.argmax() == 4:
-      output = "personne habillée" 
+    if prediction[0][prediction.argmax()]>0.6:
+      if prediction.argmax() == 0:
+        output = "Rifle violence"
+      elif prediction.argmax() == 1:
+        output = "guns violence"
+      elif prediction.argmax() == 2:
+        output = "knife violence"
+      elif prediction.argmax() == 3:
+        output = "image porno"
+      elif prediction.argmax() == 4:
+        output = "personne habillée" 
+      else:
+        output = "tank violence" 
     else:
-      output = "tank violence" 
+      output = ""
     return output