error handling

app.py

@@ -53,35 +53,38 @@ val_transform = transforms.Compose([
 def predict_image(image, model_choice):
     global model, current_model
 
+    # Check if a model is selected
+    if model_choice not in model_paths:
+        return "Error: Please select a valid model.", ""
+
     # Load the selected model if it's not already loaded
     if model_choice != current_model:
         model_path = model_paths[model_choice]
         model = load_model(model_path, model_choice)
         current_model = model_choice
-    
-    # Convert the NumPy array to a PIL Image
+    # Convert the NumPy array to a PIL Image if needed
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image.astype('uint8'), 'RGB')
-    
+   
     image = val_transform(image)
     image = image.unsqueeze(0)  # Add batch dimension
-    
+
     with torch.no_grad():
         outputs = model(image)
         probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
         confidence, preds = torch.max(probabilities, 0)
-    
+
     confidence_score = confidence.item() * 100
-    
+
     if confidence_score < 30:
-        result = "Not identified"
+        result = "Not identified\nTry to crop the image or choose another one!"
         html_result = ""
     else:
         class_name = class_names[preds.item()]
         wiki_link = class_info[preds.item()]
         result = f"{class_name}: {confidence_score:.2f}%"
         html_result = f"<h1><br><a href='{wiki_link}' target='_blank'>More Info</a></h1>"
-    
+
     return result, html_result
 
 # Load class names and class info
@@ -104,11 +107,13 @@ model = load_model(model_paths[current_model], current_model)
 
 # Create the Gradio interface
 iface = gr.Interface(
-    fn=predict_image, 
-    inputs=[gr.Image(height=500),
-            gr.Dropdown(choices=["densenet121", "resnet18", "mobilenetv2"], value="densenet121", label="Select Model")],
+    fn=predict_image,
+    inputs=[
+        gr.Image(height=500),
+        gr.Dropdown(choices=["densenet121", "resnet18", "mobilenetv2"], value="densenet121", label="Select Model")
+    ],
     outputs=[gr.Label(num_top_classes=1), gr.HTML()],
-    title="Animnal Classification",
+    title="Animal Classification",
     description="Upload an image to get the predicted label",
     allow_flagging="never",
 )