Spaces:
Sleeping
Sleeping
File size: 2,215 Bytes
cf8a522 fda9c54 cf8a522 fda9c54 cf8a522 fda9c54 cf8a522 885b5d3 cf8a522 fda9c54 cf8a522 885b5d3 cf8a522 fda9c54 cf8a522 885b5d3 cf8a522 885b5d3 fda9c54 cf8a522 fda9c54 885b5d3 fda9c54 cf8a522 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
from transformers import pipeline
from pdf2image import convert_from_path
from PIL import Image
# Initialize Flask app
app = Flask(__name__)
# Set upload folder
UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
# Allowed file extensions
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'pdf'}
# Load TrOCR Model
ocr_pipeline = pipeline("image-to-text", model="microsoft/trocr-small-printed")
def allowed_file(filename):
"""Check if the file has an allowed extension."""
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
def extract_text_from_image(image_path):
"""Extract text from a single image using TrOCR."""
image = Image.open(image_path).convert("RGB")
text = ocr_pipeline(image)[0]['generated_text']
return text
def extract_text_from_pdf(pdf_path):
"""Convert PDF to images and extract text from each page."""
images = convert_from_path(pdf_path)
extracted_text = ""
for img in images:
text = extract_text_from_image(img)
extracted_text += text + "\n"
return extracted_text.strip()
@app.route('/upload', methods=['POST'])
def upload_file():
"""Handle file upload and text extraction."""
if 'file' not in request.files:
return jsonify({"error": "No file uploaded"}), 400
file = request.files['file']
if file.filename == '':
return jsonify({"error": "No file selected"}), 400
if not allowed_file(file.filename):
return jsonify({"error": "Invalid file type. Allowed: PNG, JPG, JPEG, PDF."}), 400
# Save uploaded file
filename = secure_filename(file.filename)
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
file.save(file_path)
# Process image or PDF
if filename.lower().endswith(".pdf"):
extracted_text = extract_text_from_pdf(file_path)
else:
extracted_text = extract_text_from_image(file_path)
return jsonify({"extracted_text": extracted_text})
# Run Flask App
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000, debug=True) |