added show_error=True

app.py

@@ -77,45 +77,48 @@ def predict_image(inp):
     """
     Performs inference for a given input image and returns the prediction and CAM image.
     """
-    material_label, material_label_idx, material_probs = materials_model.predict(inp)
-    material_preds = {name: prob for name, prob in zip(material_names, material_probs.tolist())}
-
-    if material_label == 'paper':
-        specific_label, specific_label_idx, specific_probs = paper_model.predict(inp)
-        specific_preds = {name: prob for name, prob in zip(paper_names, specific_probs.tolist())}
-        specific_label = paper_names[int(specific_label_idx)]
-        recyclable_qn = paper_item_num_dict[specific_label][0]
-        recyclable_advice = paper_item_num_dict[specific_label][1]
-
-    elif material_label == 'plastic':
-        specific_label, specific_label_idx, specific_probs = plastic_model.predict(inp)
-        specific_preds = {name: prob for name, prob in zip(plastic_names, specific_probs.tolist())}
-        specific_label = plastic_names[int(specific_label_idx)]
-        recyclable_qn = plastic_item_num_dict[specific_label][0]
-        recyclable_advice = plastic_item_num_dict[specific_label][1]
-
-    elif material_label == 'glass':
-        specific_label, specific_label_idx, specific_probs = glass_model.predict(inp)
-        specific_preds = {name: prob for name, prob in zip(glass_names, specific_probs.tolist())}
-        specific_label = glass_names[int(specific_label_idx)]
-        recyclable_qn = glass_item_num_dict[specific_label][0]
-        recyclable_advice = glass_item_num_dict[specific_label][1]
-
-    elif material_label == 'metal':
-        specific_label, specific_label_idx, specific_probs = metal_model.predict(inp)
-        specific_preds = {name: prob for name, prob in zip(metal_names, specific_probs.tolist())}
-        specific_label = metal_names[int(specific_label_idx)]
-        recyclable_qn = metal_item_num_dict[specific_label][0]
-        recyclable_advice = metal_item_num_dict[specific_label][1]
-
-    elif material_label == 'others':
-        specific_label, specific_label_idx, specific_probs = others_model.predict(inp)
-        specific_preds = {name: prob for name, prob in zip(other_names, specific_probs.tolist())}
-        specific_label = other_names[int(specific_label_idx)]
-        recyclable_qn = others_item_num_dict[specific_label][0]
-        recyclable_advice = others_item_num_dict[specific_label][1]
-
-    return material_preds, specific_preds, recyclable_qn, recyclable_advice
+    try:
+      material_label, material_label_idx, material_probs = materials_model.predict(inp)
+      material_preds = {name: prob for name, prob in zip(material_names, material_probs.tolist())}
+
+      if material_label == 'paper':
+          specific_label, specific_label_idx, specific_probs = paper_model.predict(inp)
+          specific_preds = {name: prob for name, prob in zip(paper_names, specific_probs.tolist())}
+          specific_label = paper_names[int(specific_label_idx)]
+          recyclable_qn = paper_item_num_dict[specific_label][0]
+          recyclable_advice = paper_item_num_dict[specific_label][1]
+
+      elif material_label == 'plastic':
+          specific_label, specific_label_idx, specific_probs = plastic_model.predict(inp)
+          specific_preds = {name: prob for name, prob in zip(plastic_names, specific_probs.tolist())}
+          specific_label = plastic_names[int(specific_label_idx)]
+          recyclable_qn = plastic_item_num_dict[specific_label][0]
+          recyclable_advice = plastic_item_num_dict[specific_label][1]
+
+      elif material_label == 'glass':
+          specific_label, specific_label_idx, specific_probs = glass_model.predict(inp)
+          specific_preds = {name: prob for name, prob in zip(glass_names, specific_probs.tolist())}
+          specific_label = glass_names[int(specific_label_idx)]
+          recyclable_qn = glass_item_num_dict[specific_label][0]
+          recyclable_advice = glass_item_num_dict[specific_label][1]
+
+      elif material_label == 'metal':
+          specific_label, specific_label_idx, specific_probs = metal_model.predict(inp)
+          specific_preds = {name: prob for name, prob in zip(metal_names, specific_probs.tolist())}
+          specific_label = metal_names[int(specific_label_idx)]
+          recyclable_qn = metal_item_num_dict[specific_label][0]
+          recyclable_advice = metal_item_num_dict[specific_label][1]
+
+      elif material_label == 'others':
+          specific_label, specific_label_idx, specific_probs = others_model.predict(inp)
+          specific_preds = {name: prob for name, prob in zip(other_names, specific_probs.tolist())}
+          specific_label = other_names[int(specific_label_idx)]
+          recyclable_qn = others_item_num_dict[specific_label][0]
+          recyclable_advice = others_item_num_dict[specific_label][1]
+
+      return material_preds, specific_preds, recyclable_qn, recyclable_advice
+    except:
+      raise Exception("Invalid file format! Please only upload .jpg or .png files!")
 
 
 with gr.Blocks(title="Trash Classification", css="#custom_header {min-height: 3rem} #custom_title {min-height: 3rem; text-align: center}") as demo: