Spaces:
Sleeping
Sleeping
| from config import Settings | |
| import torch | |
| from PIL import Image | |
| import io | |
| from contextlib import asynccontextmanager | |
| from transformers import VisionEncoderDecoderModel | |
| from fastapi import FastAPI, UploadFile, Form, HTTPException | |
| from transformers import TrOCRProcessor, AutoTokenizer, ViTImageProcessor | |
| config = {} | |
| async def lifespan(app: FastAPI): | |
| settings = Settings() | |
| config['settings'] = settings | |
| config['device'] = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| tokenizer = AutoTokenizer.from_pretrained(settings.TOKENIZER) | |
| feature_extractor = ViTImageProcessor.from_pretrained(settings.FEATURE_EXTRACTOR) | |
| config['processor'] = TrOCRProcessor(image_processor=feature_extractor, tokenizer=tokenizer) | |
| config['ocr_model'] = VisionEncoderDecoderModel.from_pretrained(settings.OCR_MODEL) | |
| yield | |
| # Clean up and release the resources | |
| config.clear() | |
| app = FastAPI(lifespan=lifespan) | |
| def api_home(): | |
| return {'detail': 'Welcome to Sinhala OCR Space'} | |
| async def ApplyOCR(file: UploadFile): | |
| try: | |
| # Read the uploaded image file | |
| contents = await file.read() | |
| image = Image.open(io.BytesIO(contents)).convert("RGB") | |
| pixel_values = config['processor'](image, return_tensors="pt").pixel_values | |
| generated_ids = config['ocr_model'].generate(pixel_values) | |
| generated_text = config['processor'].batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| # Return the extracted text as the response | |
| return {"ocr_result": generated_text} | |
| except Exception as e: | |
| # Handle any exceptions that may occur | |
| return {"error": str(e)} |