Add livestock model to demo

app.py

@@ -10,21 +10,21 @@ def show_trees(img_path):
         raise gr.Error(f"Error initializing the deepforest model: {model_error}")
 
     img=cv2.imread(img_path)
-    
+
     if img is None:
         raise gr.Error(f"Image path is not valid.")
-    
+
     if img.shape[0]<1000 and img.shape[1]<1000:
         img = model.predict_image(path=img_path, return_plot=True)
     elif img.shape[0]>1000 or img.shape[1]>1000:
         img=model.predict_image(path=img_path,return_plot=True,thickness=2)
     else:
-        img = model.predict_image(path=img_path, return_plot=True,thickness=3) 
-    
+        img = model.predict_image(path=img_path, return_plot=True,thickness=3)
+
     if img is None:
         raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-   
+
     return img
 
 def show_birds(img_path):
@@ -34,21 +34,45 @@ def show_birds(img_path):
     except Exception as model_error:
         raise gr.Error(f"Error initializing the deepforest model: {model_error}")
     img=cv2.imread(img_path)
-    
+
+    if img is None:
+        raise gr.Error(f"Image path is not valid.")
+
+    if img.shape[0]<1000 and img.shape[1]<1000:
+        img = model.predict_image(path=img_path, return_plot=True)
+    elif img.shape[0]>1000 or img.shape[1]>1000:
+        img=model.predict_image(path=img_path,return_plot=True,thickness=2)
+    else:
+        img = model.predict_image(path=img_path, return_plot=True,thickness=3)
+
+    if img is None:
+        raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    return img
+
+def show_livestock(img_path):
+    try:
+        model = main.deepforest()
+        model.load_model(model_name="weecology/deepforest-livestock")
+    except Exception as model_error:
+        raise gr.Error(f"Error initializing the deepforest model: {model_error}")
+    img=cv2.imread(img_path)
+
     if img is None:
         raise gr.Error(f"Image path is not valid.")
-    
+
     if img.shape[0]<1000 and img.shape[1]<1000:
         img = model.predict_image(path=img_path, return_plot=True)
     elif img.shape[0]>1000 or img.shape[1]>1000:
         img=model.predict_image(path=img_path,return_plot=True,thickness=2)
     else:
-        img = model.predict_image(path=img_path, return_plot=True,thickness=3) 
-    
+        img = model.predict_image(path=img_path, return_plot=True,thickness=3)
+
     if img is None:
         raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-   
+
     return img
 
 with gr.Blocks() as demo:
@@ -63,16 +87,26 @@ with gr.Blocks() as demo:
 
     submit_button_trees = gr.Button("Predict trees")
     submit_button_trees.click(show_trees, inputs=input_image, outputs=output_image)
-    
+
     gr.Markdown('### Predict birds')
     with gr.Row():
         with gr.Column():
             input_image=gr.Image(label="Input image",type="filepath")
         with gr.Column():
             output_image=gr.Image(label="Predicted Image")
-            
+
     submit_button_birds = gr.Button("Predict birds")
     submit_button_birds.click(show_birds,inputs=input_image,outputs=output_image)
 
+    gr.Markdown('### Predict livestock')
+    with gr.Row():
+        with gr.Column():
+            input_image=gr.Image(label="Input image",type="filepath")
+        with gr.Column():
+            output_image=gr.Image(label="Predicted Image")
+
+    submit_button_livestock = gr.Button("Predict livestock")
+    submit_button_livestock.click(show_livestock,inputs=input_image,outputs=output_image)
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()