Update main.py

main.py

@@ -1,7 +1,8 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from paddleocr import PaddleOCR
 from PIL import Image
 import io, numpy as np
+import fitz  # pymupdf
 
 app = FastAPI()
 ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False)
@@ -12,16 +13,38 @@ def health():
 
 @app.post("/ocr")
 async def run_ocr(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    img_array = np.array(image)
-    result = ocr.ocr(img_array, cls=True)
+    try:
+        contents = await file.read()
+        all_texts = []
 
-    texts = []
-    for line in result[0]:
-        texts.append({
-            "text": line[1][0],
-            "confidence": round(line[1][1], 4)
-        })
+        # Handle PDF
+        if file.filename.lower().endswith(".pdf"):
+            pdf = fitz.open(stream=contents, filetype="pdf")
+            for page in pdf:
+                pix = page.get_pixmap(dpi=200)
+                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                img_array = np.array(img)
+                result = ocr.ocr(img_array, cls=True)
+                if result and result[0]:
+                    for line in result[0]:
+                        all_texts.append({
+                            "text": line[1][0],
+                            "confidence": round(line[1][1], 4)
+                        })
 
-    return {"filename": file.filename, "results": texts}
+        # Handle image (JPG, PNG, etc)
+        else:
+            image = Image.open(io.BytesIO(contents)).convert("RGB")
+            img_array = np.array(image)
+            result = ocr.ocr(img_array, cls=True)
+            if result and result[0]:
+                for line in result[0]:
+                    all_texts.append({
+                        "text": line[1][0],
+                        "confidence": round(line[1][1], 4)
+                    })
+
+        return {"filename": file.filename, "results": all_texts}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))