HawkeyeHS commited on
Commit
373c783
1 Parent(s): dd4cd7f

Added error handling for model prediction

Browse files
Files changed (1) hide show
  1. app.py +42 -7
app.py CHANGED
@@ -2,19 +2,53 @@ import gradio as gr
2
  from deepforest import main
3
  import cv2
4
 
5
-
6
  def show_trees(img_path):
7
- model = main.deepforest()
8
- model.use_release()
9
- img = model.predict_image(path=img_path, return_plot=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
11
  return img
12
 
13
  def show_birds(img_path):
14
- model = main.deepforest()
15
- model.use_bird_release()
16
- img = model.predict_image(path=img_path, return_plot=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
18
  return img
19
 
20
  with gr.Blocks() as demo:
@@ -36,6 +70,7 @@ with gr.Blocks() as demo:
36
  input_image=gr.Image(label="Input image",type="filepath")
37
  with gr.Column():
38
  output_image=gr.Image(label="Predicted Image")
 
39
  submit_button_birds = gr.Button("Predict birds")
40
  submit_button_birds.click(show_birds,inputs=input_image,outputs=output_image)
41
 
 
2
  from deepforest import main
3
  import cv2
4
 
 
5
  def show_trees(img_path):
6
+ try:
7
+ model = main.deepforest()
8
+ model.use_release()
9
+ except Exception as model_error:
10
+ raise gr.Error(f"Error initializing the deepforest model: {model_error}")
11
+
12
+ img=cv2.imread(img_path)
13
+
14
+ if img is None:
15
+ raise gr.Error(f"Image path is not valid.")
16
+
17
+ if img.shape[0]<1000 and img.shape[1]<1000:
18
+ img = model.predict_image(path=img_path, return_plot=True)
19
+ elif img.shape[0]>1000 or img.shape[1]>1000:
20
+ img=model.predict_image(path=img_path,return_plot=True,thickness=2)
21
+ else:
22
+ img = model.predict_image(path=img_path, return_plot=True,thickness=3)
23
+
24
+ if img is None:
25
+ raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
26
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
27
+
28
  return img
29
 
30
  def show_birds(img_path):
31
+ try:
32
+ model = main.deepforest()
33
+ model.use_bird_release()
34
+ except Exception as model_error:
35
+ raise gr.Error(f"Error initializing the deepforest model: {model_error}")
36
+ img=cv2.imread(img_path)
37
+
38
+ if img is None:
39
+ raise gr.Error(f"Image path is not valid.")
40
+
41
+ if img.shape[0]<1000 and img.shape[1]<1000:
42
+ img = model.predict_image(path=img_path, return_plot=True)
43
+ elif img.shape[0]>1000 or img.shape[1]>1000:
44
+ img=model.predict_image(path=img_path,return_plot=True,thickness=2)
45
+ else:
46
+ img = model.predict_image(path=img_path, return_plot=True,thickness=3)
47
+
48
+ if img is None:
49
+ raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
50
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
51
+
52
  return img
53
 
54
  with gr.Blocks() as demo:
 
70
  input_image=gr.Image(label="Input image",type="filepath")
71
  with gr.Column():
72
  output_image=gr.Image(label="Predicted Image")
73
+
74
  submit_button_birds = gr.Button("Predict birds")
75
  submit_button_birds.click(show_birds,inputs=input_image,outputs=output_image)
76