import cv2
import numpy as np
import streamlit as st
from datasets import load_dataset
from PIL import Image
from matplotlib.pyplot import imshow
from transformers import pipeline
from transformers import AutoFeatureExtractor, AutoModelForImageClassification

def load_dataset_from_Hugging_Face():
  reloaded_dataset_thermal = load_dataset("gokulraja17/rice-thermal-demo")
  reloaded_dataset_rgb = load_dataset("gokulraja17/rice-rgb-demo2")
  return reloaded_dataset_thermal,reloaded_dataset_rgb

def load_model_from_Hugging_Face():
  feature_extractor_thermal = AutoFeatureExtractor.from_pretrained("gokulraja17/convnext-fine-tuned-thermal-demo2")
  model_thermal = AutoModelForImageClassification.from_pretrained("gokulraja17/convnext-fine-tuned-thermal-demo2")
  feature_extractor_rgb = AutoFeatureExtractor.from_pretrained("gokulraja17/convnext-fine-tuned-rgb-demo2")
  model_rgb = AutoModelForImageClassification.from_pretrained("gokulraja17/convnext-fine-tuned-rgb-demo2")
  return feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb

def get_image_and_label(reloaded_dataset_thermal,reloaded_dataset_rgb,index):
  url_thermal = reloaded_dataset_thermal["test"][index]
  url_rgb = reloaded_dataset_rgb["test"][index]
  return url_thermal["image"],url_thermal["label"],url_rgb["image"],url_rgb["label"]

def create_pipeline(feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb):

  pipe_rgb = pipeline("image-classification",
                model=model_rgb,
                feature_extractor=feature_extractor_rgb)
  pipe_thermal = pipeline("image-classification",
                model=model_thermal,
                feature_extractor=feature_extractor_thermal)
  return pipe_thermal,pipe_rgb

def get_output_label(result):
  max = -999
  label = ""
  for i in result:
    if(max<i['score']):
      max = i['score']
      label = i["label"]
  return label

dataset_thermal,dataset_rgb = load_dataset_from_Hugging_Face()
labels=dataset_thermal["train"].features["label"].names
id2label={k:v for k,v in enumerate(labels)}
print(id2label)
label2id = {v:k for k,v in enumerate(labels)}
print(label2id)

feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb = load_model_from_Hugging_Face()
index=3
image_thermal,thermal_label,image_rgb,rgb_label = get_image_and_label(dataset_thermal,dataset_rgb,index)
print("Thermal Input...")
imshow(np.asarray(image_thermal))
print(id2label[thermal_label])
print("RGB Input...")
imshow(np.asarray(image_rgb))
print(id2label[rgb_label])

def consolidate_score(thermal_result,rgb_result):
  consol_score_1=[]
  consol_score_2=[]
  consol_score=[]
  for i in range(0,len(thermal_result)):
    consol_score_1.append(list(thermal_result[i].values()))
    consol_score_2.append(list(rgb_result[i].values()))
  temp_label=""
  for i in range(0,len(thermal_result)):
    temp_label=consol_score_1[i][1]
    for j in range(0,len(rgb_result)):
      if(temp_label==consol_score_2[j][1]):
        consol_score.append([temp_label,consol_score_1[i][0],consol_score_2[j][0]])
  output_consol=[]
  for i in consol_score:
    output_consol.append({'score':(((i[2]/100)*35)+((i[1]/100)*65)),'label':i[0]})
  return output_consol

pipe_thermal,pipe_rgb = create_pipeline(feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb)
thermal_result=pipe_thermal(image_thermal)
rgb_result=pipe_rgb(image_rgb)
consolidated_result=get_output_label(consolidate_score(thermal_result,rgb_result))
print(consolidated_result)

uploaded_file = st.file_uploader("Choose a image file", type=["jpg","png"])

if uploaded_file is not None:
    # Convert the file to an opencv image.
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    opencv_image = cv2.imdecode(file_bytes, 1)
    
    # Now do something with the image! For example, let's display it:
    st.image(opencv_image, channels="BGR")