Spaces:
Paused
Paused
| from fastapi import FastAPI, UploadFile, File | |
| from paddleocr import PPStructure | |
| from PIL import Image, ImageEnhance, ImageFilter | |
| import numpy as np, cv2, io | |
| app = FastAPI() | |
| # Loaded globally, but memory footprint is constrained by HF Free Tier limits | |
| table_engine = PPStructure(show_log=False, image_orientation=True, lang='en') | |
| def home(): return {"status": "Production Engine Ready"} | |
| def _deskew(image_np): | |
| gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) | |
| edges = cv2.Canny(gray, 50, 150, apertureSize=3) | |
| lines = cv2.HoughLinesP(edges, 1, np.pi/180, 100, minLineLength=100, maxLineGap=10) | |
| if lines is not None: | |
| angles = [np.degrees(np.arctan2(l[0][3]-l[0][1], l[0][2]-l[0][0])) for l in lines] | |
| valid_angles = [a for a in angles if -15 < a < 15] | |
| if valid_angles and abs(np.median(valid_angles)) > 0.5: | |
| h, w = image_np.shape[:2] | |
| M = cv2.getRotationMatrix2D((w//2, h//2), np.median(valid_angles), 1.0) | |
| return cv2.warpAffine(image_np, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255)) | |
| return image_np | |
| def _prep(img): | |
| img_np = np.array(img) | |
| # Optimization: Only run heavy deskew and contrast if it's a messy scan | |
| if cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).std() < 55: | |
| img_np = _deskew(img_np) | |
| pil_img = ImageEnhance.Contrast(Image.fromarray(img_np).convert("L")).enhance(1.8).filter(ImageFilter.SHARPEN).convert("RGB") | |
| return np.array(pil_img) | |
| return img_np | |
| def _sort(res): | |
| # FIXED: Properly handles PPStructure's nested tuple format | |
| b = [] | |
| for item in res: | |
| if isinstance(item, list) and len(item) == 2 and isinstance(item[1], (tuple, list)): | |
| box, (text, conf) = item[0], item[1] | |
| if float(conf) >= 0.45: | |
| b.append({"text": text, "y": min(p[1] for p in box), "x": min(p[0] for p in box)}) | |
| return "\n".join([x["text"] for x in sorted(b, key=lambda k: (round(k["y"]/15), k["x"]))]) | |
| async def get_ocr(file: UploadFile = File(...)): | |
| try: | |
| res = table_engine(_prep(Image.open(io.BytesIO(await file.read())).convert("RGB"))) | |
| text_blocks = [] | |
| for r in res: | |
| if r["type"] == "table": | |
| text_blocks.append(f"[TABLE]\n{r.get('res', {}).get('html', '')}") | |
| else: | |
| sorted_txt = _sort(r.get("res", [])) | |
| if sorted_txt: text_blocks.append(sorted_txt) | |
| return {"text": "\n\n".join(text_blocks)} | |
| except Exception as e: | |
| return {"text": "", "error": str(e)} |