Update app.py

app.py

@@ -1,35 +1,31 @@
 import gradio as gr
-import tensorflow as tf
+import torch
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 import keras_ocr
-import requests
 import cv2
-import os
-import csv
-import numpy as np
-import pandas as pd
-import huggingface_hub
-from huggingface_hub import Repository
-from datetime import datetime
-import scipy.ndimage.interpolation as inter
 import easyocr
-import datasets
-from datasets import load_dataset, Image
-from PIL import Image
 from paddleocr import PaddleOCR
-from save_data import flag
-            
+import numpy as np
+
+# Load tokenizer
+tokenizer = DistilBertTokenizer.from_pretrained("./distilbert_spam_model")
+
+# Load model
+model = DistilBertForSequenceClassification.from_pretrained("./distilbert_spam_model")
+model.load_state_dict(torch.load("./distilbert_spam_model/model.pth", map_location=torch.device('cpu')))
+model.eval()
+
 """
 Paddle OCR
 """
 def ocr_with_paddle(img):
     finaltext = ''
     ocr = PaddleOCR(lang='en', use_angle_cls=True)
-    # img_path = 'exp.jpeg'
     result = ocr.ocr(img)
     
     for i in range(len(result[0])):
         text = result[0][i][1][0]
-        finaltext += ' '+ text
+        finaltext += ' ' + text
     return finaltext
 
 """
@@ -37,84 +33,62 @@ Keras OCR
 """
 def ocr_with_keras(img):
     output_text = ''
-    pipeline=keras_ocr.pipeline.Pipeline()
-    images=[keras_ocr.tools.read(img)]
-    predictions=pipeline.recognize(images)
-    first=predictions[0]
-    for text,box in first:
-        output_text += ' '+ text
+    pipeline = keras_ocr.pipeline.Pipeline()
+    images = [keras_ocr.tools.read(img)]
+    predictions = pipeline.recognize(images)
+    
+    for text, _ in predictions[0]:
+        output_text += ' ' + text
     return output_text
 
 """
-easy OCR
+Easy OCR
 """
-# gray scale image
-def get_grayscale(image):
-    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-
-# Thresholding or Binarization
-def thresholding(src):
-    return cv2.threshold(src,127,255, cv2.THRESH_TOZERO)[1]
 def ocr_with_easy(img):
-    gray_scale_image=get_grayscale(img)
-    thresholding(gray_scale_image)
-    cv2.imwrite('image.png',gray_scale_image)
-    reader = easyocr.Reader(['th','en'])
-    bounds = reader.readtext('image.png',paragraph="False",detail = 0)
-    bounds = ''.join(bounds)
-    return bounds
-        
+    reader = easyocr.Reader(['en'])
+    bounds = reader.readtext(img, paragraph=True, detail=0)
+    return ' '.join(bounds)
+
 """
-Generate OCR
+Generate OCR and classify spam
 """
-def generate_ocr(Method,img):
-    
+def generate_ocr_and_classify(Method, img):
+    if img is None:
+        raise gr.Error("Please upload an image!")
+
+    # Perform OCR
     text_output = ''
-    if (img).any():
-        add_csv = []
-        image_id = 1
-        print("Method___________________",Method)
-        if Method == 'EasyOCR':
-            text_output = ocr_with_easy(img)
-        if Method == 'KerasOCR':
-            text_output = ocr_with_keras(img)
-        if Method == 'PaddleOCR':
-            text_output = ocr_with_paddle(img)
-       
-        try:
-            flag(Method,text_output,img)
-        except Exception as e:
-            print(e)
-        return text_output
-    else:
-        raise gr.Error("Please upload an image!!!!")
+    if Method == 'EasyOCR':
+        text_output = ocr_with_easy(img)
+    elif Method == 'KerasOCR':
+        text_output = ocr_with_keras(img)
+    elif Method == 'PaddleOCR':
+        text_output = ocr_with_paddle(img)
     
-    # except Exception as e:
-    #     print("Error in ocr generation ==>",e)
-    #     text_output = "Something went wrong"
-    # return text_output
+    # Classify extracted text
+    inputs = tokenizer(text_output, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
     
+    prediction = torch.argmax(outputs.logits, dim=1).item()
+    classification = "Spam" if prediction == 1 else "Not Spam"
+
+    return text_output, classification
 
 """
-Create user interface for OCR demo
+Create user interface
 """
-
-# image = gr.Image(shape=(300, 300))
 image = gr.Image()
-method = gr.Radio(["PaddleOCR","EasyOCR", "KerasOCR"],value="PaddleOCR")
-output = gr.Textbox(label="Output")
+method = gr.Radio(["PaddleOCR", "EasyOCR", "KerasOCR"], value="PaddleOCR")
+output_text = gr.Textbox(label="Extracted Text")
+output_label = gr.Label(label="Classification")
 
 demo = gr.Interface(
-    generate_ocr,
-    [method,image],
-    output,
-    title="Optical Character Recognition",
-    css=".gradio-container {background-color: lightgray} #radio_div {background-color: #FFD8B4; font-size: 40px;}",
-    article = """<p style='text-align: center;'>Feel free to give us your thoughts on this demo and please contact us at 
-                    <a href="mailto:[email protected]" target="_blank">[email protected]</a> 
-                    <p style='text-align: center;'>Developed by: <a href="https://www.pragnakalp.com" target="_blank">Pragnakalp Techlabs</a></p>"""
-    
-
+    generate_ocr_and_classify,
+    [method, image],
+    [output_text, output_label],
+    title="OCR & Spam Classification",
+    description="Upload an image with text, extract the text using OCR, and classify whether it is spam or not.",
 )
-# demo.launch(enable_queue = False)
+
 demo.launch()