import easyocr 
from gradio_client import Client, handle_file
import pandas as pd
import gradio as gr

clientImg = Client("dj-dawgs-ipd/IPD-Image-ViT-Finetune")
clientEngText = Client("dj-dawgs-ipd/IPD-Text-English-Finetune")
clientHingText = Client("dj-dawgs-ipd/IPD-Text-Hinglish")

profanity_df = pd.read_csv('Hinglish_Profanity_List.csv' , encoding = 'utf-8')
profanity_hn = profanity_df['profanity_hn']

def extract_text(image):
    reader = easyocr.Reader(['en'])
    data = [result[1] for result in reader.readtext(image)]
    return ' '.join([l for l in data])

def predict(image):
    
    imgResult = clientImg.predict(
        image=handle_file(image),
        api_name="/predict"
    )
    
    label , confidence = imgResult[0]['label'] , float(imgResult[1]['label'])    

    if (label == 'finger_gun_to_the_head' and confidence > 0.98) or (label != 'finger_gun_to_the_head' and confidence > 0.95):
        return {
            "prediction" : "hate",
            "language" : None,
            "label" : label,
            "confidence" : confidence,
            "hate_text" : None
        }
    else:
        ocr_text = extract_text(image).lower()
        
        engResult = clientEngText.predict(
		text=ocr_text[:200],
		api_name="/predict"
        )
        
        hingResult = clientHingText.predict(
        		text=ocr_text[:200],
        		api_name="/predict"
        )

        profanityFound = [word for word in ocr_text.split() if word in profanity_hn]

        if len(profanityFound) > 0:
            return {
                "prediction" : "hate",
                "language" : "Hindi",
                "label" : "Profanity Found",
                "confidence" : None,
                "hate_text" : profanityFound
            }

        elif engResult[0] != "NEITHER" and engResult[1] > 0.5:
            return {
                "prediction" : "hate",
                "language" : "English",
                "label" : engResult[0],
                "confidence" : engResult[1],
                "hate_text" : ocr_text[:200]
            }

        elif hingResult[0] != "NAG" and hingResult[1] > 0.5:
            return {
                "prediction" : "hate",
                "language" : "Hinglish",
                "label" : hingResult[0],
                "confidence" : hingResult[1],
                "hate_text" : ocr_text[:200]
            }

        else:
            return {
                "prediction" : "not_hate",
                "language" : None,
                "label" : "No hate found, yay!",
                "confidence" : None,
                "hate_text" : None
            }
        
iface = gr.Interface(fn=predict,
                     inputs = gr.Image(type='filepath'),
                     outputs=gr.JSON(),         
                     title = "Hate Speech Detection in Image",
                     description = "Detect hateful symbols or text in Image"
        )
                     
if __name__ == "__main__":
    iface.launch(show_error=True)