import os
from fastapi import FastAPI, File, UploadFile
from paddleocr import PaddleOCR
from PIL import Image
import numpy as np
import io
import uvicorn

app = FastAPI()

# Initialize PaddleOCR
ocr = PaddleOCR(lang='en', use_angle_cls=True)

@app.post("/extract-text/")
async def extract_text(file: UploadFile = File(...)):
    image = Image.open(io.BytesIO(await file.read()))
    
    # Convert the image to a NumPy array
    img_array = np.array(image)

    # Handle different image channels
    if img_array.ndim == 2:  # Grayscale image
        img_array = np.stack((img_array,)*3, axis=-1)
    elif img_array.shape[-1] == 4:  # RGBA image
        img_array = img_array[..., :3]
    
    # Perform OCR
    result = ocr.ocr(img_array)
    extracted_text = ' '.join([line[1][0] for line in result[0]])
    return {"text": extracted_text}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8000)))