Update app.py

app.py

@@ -1,52 +1,39 @@
 import gradio as gr
-
-def upload_image(image):
-    return image
-
-iface = gr.Interface(
-    fn=upload_image, 
-    inputs=gr.Image(type="pil",  image_mode="RGB"), 
-    outputs="image",
-    title="Image Uploader",
-    description="Upload an image and see the result below."
-)
-
-# from PIL import Image
-# import os
-# import numpy as np
-# import urllib.request
-# import glob
-
-# # intake library and plugin
-# # import intake
-# # from intake_zenodo_fetcher import download_zenodo_files_for_entry
-
-# # geospatial libraries
-# import geopandas as gpd
-
-# from rasterio.transform import from_origin
-# import rasterio.features
-
-# import fiona
-
-# from shapely.geometry import shape, mapping, box
-# from shapely.geometry.multipolygon import MultiPolygon
-
-# # machine learning libraries
-# from detectron2 import model_zoo
-# from detectron2.engine import DefaultPredictor
-# from detectron2.utils.visualizer import Visualizer, ColorMode
-# from detectron2.config import get_cfg
-# from detectron2.engine import DefaultTrainer
-
-# # visualisation
-# # import holoviews as hv
-# # from IPython.display import display
-# # import geoviews.tile_sources as gts
-
-# # import hvplot.pandas
-# # import hvplot.xarray
-
-# hv.extension('bokeh', width=100)
-
-
+from deepforest import main
+import matplotlib.pyplot as plt
+
+# Initialize the deepforest model and use the released version
+model = main.deepforest()
+model.use_release()
+
+def predict_and_visualize(image):
+    """
+    Function to predict and visualize the image using deepforest model.
+    
+    Args:
+    - image: An image array.
+    
+    Returns:
+    - An image with predictions visualized.
+    """
+    # Predict image and return plot. Since Gradio passes image as array, save it temporarily.
+    temp_path = "/tmp/uploaded_image.png"
+    plt.imsave(temp_path, image)
+    img = model.predict_image(path=temp_path, return_plot=True)
+    
+    # Since the output is BGR and matplotlib (and hence Gradio) needs RGB, we convert the color scheme
+    img_rgb = img[:, :, ::-1]
+    
+    # Return the RGB image
+    return img_rgb
+
+# Define the Gradio interface
+iface = gr.Interface(fn=predict_and_visualize,
+                     inputs=gr.Image(type="numpy", label="Upload Image"),
+                     outputs=gr.Image(label="Predicted Image"),
+                     title="DeepForest Tree Detection",
+                     examples=["./example.jpg"],
+                     description="Upload an image to detect trees using the DeepForest model.")
+
+# Launch the Gradio app
+iface.launch()