import tensorflow as tf
import gradio as gr
import tensorflow_hub as hub
from PIL import Image
from ultralytics import YOLO


#Transforming or Reading an input image and making it prediction ready as per our model
#filename :- input image path as str
#img_shape :- input image shape except color channel as List i.e here [128,128]
def transform_image_for_prediction(filename,img_shape=[256,256]):
  """
  Reads an image from filename, turns it into a tensor
  and reshapes it to (img_shape, img_shape, colour_channel).
  """
  # Read in target file (an image)
  img = tf.io.read_file(filename)
  
  # Decode the read file (array) into a tensor & ensure 3 colour channels 
  # (our model is trained on images with 3 colour channels and sometimes images have 4 colour channels)
  img = tf.image.decode_image(img, channels=3)
  
  # Resize the image (to the same size our model was trained on)
  img = tf.image.resize(img,size=img_shape)
  
  #Applying Feature scaling as we done it /255. in our training and test data
  img = img/255.                 #dot at end of 255 so to convert integer into float value 
  
  return img


def make_prediction(img):
    #cnn_model = pickle.load(open("model.pkl","rb"))
    cnn_model =  tf.keras.models.load_model("cnn.h5")
    resnet_model = tf.keras.models.load_model("resnetV2.h5",custom_objects={'KerasLayer':hub.KerasLayer})
    
    img_tf = transform_image_for_prediction(img)
    pred_prob = cnn_model.predict(tf.expand_dims(img_tf, axis=0)) #adding fake dimension which represnts batchno
    pred_prob_resnet = resnet_model.predict(tf.expand_dims(img_tf,axis=0))
    
    # class_names = ['crocodile', 'logs'] #as define in .ipynb file
    # if len(pred_prob[0]) > 1: # prediction for multi-class
    #     pred_class = class_names[pred_prob.argmax()] # if more than one output, take the max
    # else:                #prediction for binary class
    #     if (pred_prob[0][0] >=0.60):
    #         i=1
    #     else:
    #         i=0
    #     pred_class = class_names[i] # if only one output, round
    
    #result = f"Prediction of being crocodile: {np.round(((1-pred_prob[0][0])*100),2)}%" + "\n" + f"Prediction of being driftwood: {np.round(((pred_prob[0][0])*100),2)}%"   
    cnn_result = {"crocodile": float(pred_prob[0][0]), "driftwood": float(pred_prob[0][1]),"other":float(pred_prob[0][2])}
    resnet_result = {"crocodile": float(pred_prob_resnet[0][0]), "driftwood": float(pred_prob_resnet[0][1]),"other":float(pred_prob_resnet[0][2])}

    '''
    Appending output of Object Detection model trained on custom dataset ,annotated using RoboFlow website.
    
    Dataset Info :-  https://universe.roboflow.com/pujii4ml/croc-o-net
    
    '''
    yolo = YOLO(model="yolov8_medium_25.pt")     #put here pretrained saved weights file of your trained yolov8 model 
    results = yolo.predict(img,conf=0.15)
    objectDetectionResult=Image.fromarray(results[0].plot()[:,:,::-1])
    
    return cnn_result,resnet_result,objectDetectionResult

'''
Image component input Arguments  :- https://www.gradio.app/docs/image
-------------------------------

type = The format the image is converted to before being passed into the prediction function. 
       "numpy" converts the image to a numpy array with shape (height, width, 3) and values from 0 to 255, 
       "pil" converts the image to a PIL image object, 
       "filepath" passes a str path to a temporary file containing the image.
       
label = component name in interface. 

interactive = if True, will allow users to upload and edit an image; if False, can only be used to display images. 
              If not provided, this is inferred based on whether the component is used as an input or output.

'''
    

inputs = gr.Image(type='filepath',label="Feed me some image!")  
outputs=[
    gr.Label(label="CNN From Scratch using 2 Convolutional Layer"),
    gr.Label(label="ResNetV2 cnn model fine tunned with smaller learning rate"),
    gr.Image(type="pil",label="Object Detection Model YoLoV8 ,trained on custom dataset")
]

app = gr.Interface(fn=make_prediction,inputs=inputs,outputs=outputs,title="Welcome to the Croc-O-Net Classifier 🐊🔍",
                   flagging_options=["Incorrect Classification"],theme="soft",allow_flagging="manual",flagging_dir="flagged",
                   description="Tired of squinting at blurry vacation photos, wondering if that log is a laid-back crocodile or just nature’s idea of a floating joke?🐊 Fear not!",
                   article="<p style='text-align: center'><a href='https://medium.com/@p3pioneer22/croc-o-net-the-only-cnn-project-you-need-to-get-a-job-d57b86ac8fac' target='_blank'>Medium Blog Post 📰</a></p>")
app.launch()