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,
    		api_name="/predict"
    )
    
    if float(imgResult[1]['label']) > 0.95:
        return ["hate" , imgResult[0]]
    else:
        
        ocr_text = extract_text(image)
        
        engResult = clientEngText.predict(
		text=ocr_text,
		api_name="/classify_text"
        )
        
        hingResult = clientHingText.predict(
        		text=ocr_text,
        		api_name="/predict"
        )
        
        if engResult[0] == "NEITHER" or hingResult[0] == "NAG":
            return ["not_hate" , None]
        else:
            return ["hate" , None]
        
iface = gr.Interface(fn=predict,
                     inputs = gr.Image(type='pil'),
                     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()