dead031/ai-worker

1import os2import io3import json4import uvicorn5from fastapi import FastAPI, UploadFile, File, HTTPException6from pdf2image import convert_from_bytes7import layoutparser as lp8import numpy as np9import cv210import pytesseract11from PIL import Image12 13app = FastAPI()14 15import urllib.request16import os17 18def download_model():19    """Download model files to /tmp (the only writable area on HF)"""20    print("Checking model files in /tmp...")21    base_dir = "/tmp/models"22    model_dir = os.path.join(base_dir, "faster_rcnn_R_50_FPN_3x")23    os.makedirs(model_dir, exist_ok=True)24 25    files = {26        os.path.join(model_dir, "config.yaml"): 27            "https://huggingface.co/layoutparser/detectron2/resolve/main/PubLayNet/faster_rcnn_R_50_FPN_3x/config.yml",28        os.path.join(model_dir, "model_final.pth"): 29            "https://huggingface.co/layoutparser/detectron2/resolve/main/PubLayNet/faster_rcnn_R_50_FPN_3x/model_final.pth"30    }31 32    for path, url in files.items():33        if not os.path.exists(path):34            print(f"Downloading {path} from {url}...")35            urllib.request.urlretrieve(url, path)36    print("Model files ready.")37 38# Ensure models are downloaded before initialization39download_model()40 41# Initialize LayoutParser with paths pointing to /tmp42model = lp.Detectron2LayoutModel(43    config_path="/tmp/models/faster_rcnn_R_50_FPN_3x/config.yaml",44    model_path="/tmp/models/faster_rcnn_R_50_FPN_3x/model_final.pth",45    extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.5],46    label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}47)48 49@app.get("/")50def home():51    return {"message": "Deshonnati AI Worker is running"}52 53@app.post("/process")54async def process_pdf(file: UploadFile = File(...)):55    try:56        # 1. Read PDF bytes57        pdf_bytes = await file.read()58        59        # 2. Convert PDF to Image (first page for demo)60        images = convert_from_bytes(pdf_bytes, dpi=200)61        if not images:62            raise HTTPException(status_code=400, detail="Could not convert PDF to images")63        64        results = []65        66        for i, image in enumerate(images):67            # 3. Convert PIL image to CV2 format for LayoutParser68            open_cv_image = np.array(image)69            open_cv_image = open_cv_image[:, :, ::-1].copy() # RGB to BGR70 71            # 4. Detect Layout72            layout = model.detect(open_cv_image)73            74            page_articles = []75            76            # 5. Process each detected block77            for block in layout:78                if block.type in ['Text', 'Title']:79                    # Get coordinates80                    x0, y0, x1, y1 = block.coordinates81                    82                    # Crop image for better OCR83                    cropped_img = image.crop((x0, y0, x1, y1))84                    85                    # 6. Run OCR (Multi-language)86                    text = pytesseract.image_to_string(cropped_img, lang='mar+eng+hin')87                    88                    page_articles.append({89                        "id": f"art_{i}_{len(page_articles)}",90                        "type": block.type,91                        "bbox": {92                            "x0": x0,93                            "y0": y0,94                            "x1": x1,95                            "y1": y196                        },97                        "text": text.strip()98                    })99            100            results.append({101                "page": i + 1,102                "articles": page_articles103            })104            105        return {"success": True, "data": results}106 107    except Exception as e:108        raise HTTPException(status_code=500, detail=str(e))109 110if __name__ == "__main__":111    # HF Spaces requires port 7860112    uvicorn.run(app, host="0.0.0.0", port=7860)113