|
import os |
|
from flask import Flask, request, jsonify |
|
from PIL import Image |
|
import torch |
|
from torchvision.transforms import functional as F |
|
from transformers import TrOCRProcessor, VisionEncoderDecoderModel |
|
from transformers import AutoModel |
|
|
|
from transformers import AutoModel |
|
|
|
model_name = "trocrnew.pth" |
|
access_token = os.environ.get("HF_TOKEN") |
|
model = AutoModel.from_pretrained(model_name, token=access_token) |
|
|
|
|
|
|
|
|
|
|
|
app = Flask(__name__) |
|
|
|
|
|
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-small-printed") |
|
model = VisionEncoderDecoderModel.from_pretrained("trocrnew.pth") |
|
|
|
|
|
model.eval() |
|
|
|
def ocr(image): |
|
|
|
image = F.to_tensor(image).unsqueeze(0) |
|
|
|
|
|
with torch.no_grad(): |
|
generated_ids = model.generate(image) |
|
|
|
|
|
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] |
|
|
|
return generated_text |
|
|
|
@app.route('/ocr', methods=['POST']) |
|
def ocr_endpoint(): |
|
|
|
if 'file' not in request.files: |
|
return jsonify({'error': 'No file provided'}), 400 |
|
|
|
file = request.files['file'] |
|
|
|
|
|
allowed_extensions = {'png', 'jpg', 'jpeg', 'gif'} |
|
if '.' not in file.filename or file.filename.split('.')[-1].lower() not in allowed_extensions: |
|
return jsonify({'error': 'Invalid file type'}), 400 |
|
|
|
|
|
try: |
|
image = Image.open(file).convert('RGB') |
|
text = ocr(image) |
|
return jsonify({'text': text}), 200 |
|
except Exception as e: |
|
return jsonify({'error': str(e)}), 500 |
|
|
|
if __name__ == '__main__': |
|
app.run(debug=True) |
|
|