import numpy as np
import gradio as gr
from deepforest import main
from deepforest import get_data
import os
import matplotlib.pyplot as plt
import rasterio
from PIL import Image

#img = model.predict_image(path="/Users/benweinstein/Documents/NeonTreeEvaluation/evaluation/RGB/TEAK_049_2019.tif",return_plot=True)

#predict_image returns plot in BlueGreenRed (opencv style), but matplotlib likes RedGreenBlue, switch the channel order.
def predictDeepForest(input_img):
    model = main.deepforest()
    model.use_release()
    #print(input_img.shape)
    im = Image.open(input_img) 
    im = np.array(im)
    img = model.predict_image(im, return_plot=True)
    img.shape
    plt.imshow(img[:,:,::-1])
    return img
#gr.Image(type="filepath", shape=...)

demo = gr.Interface(predictDeepForest, gr.Image(type="filepath"), "image")
demo.launch(debug = True)