from IPython.display import display, JSON
import matplotlib.pyplot as plt
from speciesnet import DEFAULT_MODEL
from speciesnet import draw_bboxes
from speciesnet import load_rgb_image
from speciesnet import SpeciesNet
from speciesnet import SUPPORTED_MODELS
import time
import gradio as gr
import json


def print_predictions(predictions_dict: dict) -> None:
    print("Predictions:")
    for prediction in predictions_dict["predictions"]:
        print(prediction["filepath"], "=>", prediction["prediction"])


# ##Pick model variation
print("Default SpeciesNet model:", DEFAULT_MODEL)
print("Supported SpeciesNet models:", SUPPORTED_MODELS)
model = SpeciesNet(DEFAULT_MODEL)


def inference(image):
    # Assuming image is uploaded and accessible as a file path
    filepath = "temp_image.jpg"  # Temporary file name
    image.save(filepath)

    start_time = time.time()
    predictions_dict = model.predict(
        instances_dict={
            "instances": [
                {
                    "filepath": filepath,
                    "country": "VNM",  # You can modify this
                }
            ]
        }
    )
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"Execution time: {execution_time} seconds")

    # Format the JSON for readability
    formatted_json = json.dumps(predictions_dict, indent=4)
    return formatted_json


# Create a Gradio interface
iface = gr.Interface(
    fn=inference,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="SpeciesNet Inference",
    description="",
)

iface.launch()