import cv2 import io import numpy as np from PIL import Image import pytesseract from fastapi import FastAPI, UploadFile, File from fastapi.middleware.cors import CORSMiddleware from mltu.inferenceModel import OnnxInferenceModel from mltu.utils.text_utils import ctc_decoder from mltu.transformers import ImageResizer from mltu.configs import BaseModelConfigs from textblob import TextBlob from happytransformer import HappyTextToText, TTSettings from transformers import AutoTokenizer, T5ForConditionalGeneration from pydantic import BaseModel tokenizer = AutoTokenizer.from_pretrained("grammarly/coedit-large", cache_dir="./cache") chatModel = T5ForConditionalGeneration.from_pretrained("grammarly/coedit-large", cache_dir="./cache") configs = BaseModelConfigs.load("./configs.yaml") #happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction") beam_settings = TTSettings(num_beams=5, min_length=1, max_length=100) app = FastAPI() origins = ["*"] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class ImageToWordModel(OnnxInferenceModel): def __init__(self, char_list, *args, **kwargs): super().__init__(*args, **kwargs) self.char_list = char_list def predict(self, image: np.ndarray): image = ImageResizer.resize_maintaining_aspect_ratio( image, *self.input_shape[:2][::-1] ) image_pred = np.expand_dims(image, axis=0).astype(np.float32) preds = self.model.run(None, {self.input_name: image_pred})[0] text = ctc_decoder(preds, self.char_list)[0] return text model = ImageToWordModel(model_path=configs.model_path, char_list=configs.vocab) extracted_text = "" @app.post("/extract_handwritten_text/") async def predict_text(image: UploadFile): global extracted_text # Read the uploaded image img = await image.read() nparr = np.frombuffer(img, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # Make a prediction extracted_text = model.predict(img) #corrected_text = happy_tt.generate_text(extracted_text, beam_settings) return {"text": extracted_text} @app.post("/extract_text/") async def extract_text_from_image(image: UploadFile): global extracted_text # Check if the uploaded file is an image if image.content_type.startswith("image/"): # Read the image from the uploaded file image_bytes = await image.read() img = Image.open(io.BytesIO(image_bytes)) # Perform OCR on the image extracted_text = pytesseract.image_to_string(img) #corrected_text = happy_tt.generate_text(extracted_text, beam_settings) return {"text": extracted_text} else: return {"error": "Invalid file format. Please upload an image."} class ChatPrompt(BaseModel): prompt: str @app.post("/chat_prompt/") async def chat_prompt(request: ChatPrompt): global extracted_text input_text = request.prompt + ": " + extracted_text print(input_text) input_ids = tokenizer(input_text, return_tensors="pt").input_ids outputs = chatModel.generate(input_ids, max_length=256) edited_text = tokenizer.decode(outputs[0], skip_special_tokens=True) return {"edited_text": edited_text}