Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from transformers import DistilBertTokenizer, DistilBertForSequenceClassification | |
| import keras_ocr | |
| import cv2 | |
| import easyocr | |
| from paddleocr import PaddleOCR | |
| 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) | |
| result = ocr.ocr(img) | |
| for i in range(len(result[0])): | |
| text = result[0][i][1][0] | |
| finaltext += ' ' + text | |
| return finaltext | |
| """ | |
| Keras OCR | |
| """ | |
| def ocr_with_keras(img): | |
| output_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 | |
| """ | |
| def ocr_with_easy(img): | |
| reader = easyocr.Reader(['en']) | |
| bounds = reader.readtext(img, paragraph=True, detail=0) | |
| return ' '.join(bounds) | |
| """ | |
| Generate OCR and classify spam | |
| """ | |
| def generate_ocr_and_classify(Method, img): | |
| if img is None: | |
| raise gr.Error("Please upload an image!") | |
| # Perform OCR | |
| text_output = '' | |
| 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) | |
| # 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 | |
| """ | |
| image = gr.Image() | |
| 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_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() | |