Spaces:
Sleeping
Sleeping
from fastapi import FastAPI, File, UploadFile, HTTPException | |
from PIL import Image | |
import pytesseract | |
import numpy as np | |
import cv2 | |
import subprocess | |
import os | |
import io | |
pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract' | |
app = FastAPI() | |
def set_dpi(image: Image.Image, target_dpi=300): | |
"""Set the image DPI to the target value while preserving aspect ratio.""" | |
current_dpi = image.info.get('dpi', (72, 72))[0] | |
if current_dpi == target_dpi: | |
return image | |
# Calculate scale factor to achieve target DPI | |
scale_factor = target_dpi / current_dpi | |
# Resize image | |
width, height = image.size | |
new_size = (int(width * scale_factor), int(height * scale_factor)) | |
resized_image = image.resize(new_size, Image.Resampling.LANCZOS) | |
# Set DPI metadata | |
resized_image.info['dpi'] = (target_dpi, target_dpi) | |
return resized_image | |
def preprocess_dual(image: Image.Image): | |
"""Set DPI to 300, convert to grayscale, and produce normal and inverted versions.""" | |
# Set DPI to 300 | |
image_300dpi = set_dpi(image, target_dpi=300) | |
img = np.array(image_300dpi) | |
if img.ndim == 3: | |
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) | |
else: | |
gray = img | |
inverted = cv2.bitwise_not(gray) | |
return gray, inverted | |
def contains_croatian_chars(text): | |
cro_chars = set("čćđžšČĆĐŽŠ") | |
return any(c in text for c in cro_chars) | |
def run_ocr(image_np, lang='hrv', config='--oem 3 --psm 6'): | |
"""Run Tesseract OCR on a grayscale image""" | |
return pytesseract.image_to_string(image_np, config=config, lang=lang) | |
async def ocr(file: UploadFile = File(...)): | |
if not file: | |
raise HTTPException(status_code=400, detail={'error': 'No image uploaded'}) | |
try: | |
image = Image.open(io.BytesIO(await file.read())).convert('RGB') | |
except Exception as e: | |
raise HTTPException(status_code=400, detail={'error': f'Invalid image file: {str(e)}'}) | |
gray, inverted = preprocess_dual(image) | |
try: | |
# Croatian OCR | |
text_hrv_normal = run_ocr(gray, lang='hrv') | |
text_hrv_inverted = run_ocr(inverted, lang='hrv') | |
text_hrv = text_hrv_inverted if len(text_hrv_inverted.strip()) > len(text_hrv_normal.strip()) else text_hrv_normal | |
if contains_croatian_chars(text_hrv): | |
final_text = text_hrv | |
else: | |
# Fallback to English OCR | |
text_eng_normal = run_ocr(gray, lang='eng') | |
text_eng_inverted = run_ocr(inverted, lang='eng') | |
text_eng = text_eng_inverted if len(text_eng_inverted.strip()) > len(text_eng_normal.strip()) else text_eng_normal | |
final_text = text_eng if len(text_eng.strip()) > len(text_hrv.strip()) else text_hrv | |
final_text = final_text.replace('£', 'E') | |
final_text = final_text.replace('€', 'E') | |
except Exception as e: | |
raise HTTPException(status_code=500, detail={'error': f'OCR failed: {str(e)}'}) | |
return {'text': final_text} | |
async def debug_tesseract(): | |
try: | |
result = subprocess.run(['tesseract', '--version'], capture_output=True, text=True, check=True) | |
path = os.environ.get('PATH', 'PATH not set') | |
return { | |
'tesseract_version': result.stdout.strip(), | |
'tesseract_path': pytesseract.pytesseract.tesseract_cmd, | |
'current_path': path | |
} | |
except subprocess.CalledProcessError as e: | |
return { | |
'error': str(e), | |
'stderr': e.stderr, | |
'path': os.environ.get('PATH', 'PATH not set') | |
}, 500 | |
except FileNotFoundError: | |
return { | |
'error': 'Tesseract executable not found', | |
'path': os.environ.get('PATH', 'PATH not set') | |
}, 500 | |
except Exception as e: | |
return {'error': str(e)}, 500 | |
async def root(): | |
return {"message": "Hello from FastAPI"} |