# import tensorflow as tf

import gradio as gr
from PIL import Image

import json

# Use a pipeline as a high-level helper
from transformers import pipeline


pipe = pipeline("image-classification", model="janjibDEV/vit-plantnet300k")


# Specify the path to your JSON file
id_2_species = 'class_idx_to_species_id.json'
species_2_name = 'plantnet300K_species_id_2_name.json'

# Open and load the JSON file
with open(id_2_species, 'r') as file:
    id_2_species_dict = json.load(file)

# Open and load the JSON file
with open(species_2_name, 'r') as file:
    species_2_name_dict = json.load(file)

def get_species_name(label):
  return species_2_name_dict[id_2_species_dict[str(label)]]

def combine_pred(dics):
  combined_dict = {}
  for dic in dics:
      combined_dict[get_species_name(dic['label'])] = dic['score']
  return combined_dict

def classify_image(inp):
    res = combine_pred(pipe(Image.open(inp))[:3])
    return res

title = "Plantify: Identify a plant!"
description = """
Plantify is powered by a finetuned ViT model trained on the PlantNet300K dataset. Made by JanjibDEV
"""

gr.Interface(
    title=title,
    description=description,
    fn=classify_image,
    inputs=gr.Image(type="filepath"),
    outputs=gr.Label(num_top_classes=3),
    examples=[["marigold_pic.jpg"]],
).launch(debug=True, share=True)