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 confidence > 0.95:
        return ["hate" , f"label: {label}, confidence: {confidence}"]
    else:
 
        ocr_text = extract_text(image).lower()
        
        engResult = clientEngText.predict(
		text=ocr_text,
		api_name="/predict"
        )
        
        hingResult = clientHingText.predict(
        		text=ocr_text,
        		api_name="/predict"
        )

        profanityFound = any(word in ocr_text.split() for word in profanity_hn)

        if profanityFound:
            return ["hate", "Profanity Found"]
        elif engResult[0] != "NEITHER" and engResult[1] > 0.5:
            return ["hate", f"Result: {engResult}, Text: {ocr_text}"]
        elif hingResult[0] != "NAG" and hingResult[1] > 0.5:
            return ["hate", f"Result: {hingResult}, Text: {ocr_text}"]
        else:
            return ["not_hate", "No hate found, yay!"]
        
        # if not profanityFound and engResult[0] == "NEITHER" and hingResult[0] == "NAG":
        #     return ["not_hate" , "No Hate Symbols Detected"]
        # else:
        #     return ["hate" , "No Hate Symbols Detected"]
        
iface = gr.Interface(fn=predict,
                     inputs = gr.Image(type='filepath'),
                     outputs=[gr.Label(label = "Class (hate or not_hate)") , gr.Label(label = "Explanation")],
                     title = "Hate Speech Detection in Image",
                     description = "Detect hateful symbols or text in Image"
        )
                     
if __name__ == "__main__":
    iface.launch()