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

        profanityFound = any(word in text.split() for word in profanity_hn)
        
        if not profanityFound and (engResult[0] == "NEITHER" or 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") , gr.Label(label = "Hate Symbol(if any)")],
                     title = "Hate Speech Detection in Image",
                     description = "Detect hateful symbols or text in Image"
        )
                     
if __name__ == "__main__":
    iface.launch()