Update app.py

app.py

@@ -18,34 +18,30 @@ BUCKET_NAME = "ocr_images"
 
 app = FastAPI(title="RapidOCR FastAPI", version="1.0")
 
-# -------------------------------------------------------------------
-# DB + GRIDFS INIT
-# -------------------------------------------------------------------
+# DB INIT
 _client = AsyncIOMotorClient(MONGO_URI)
 _db = _client[DB_NAME]
-_gridfs = AsyncIOMotorGridFSBucket(_db, bucket_name=BUCKET_NAME)
 
-# RapidOCR ONNX model (loaded once)
 ocr_engine = RapidOCR()
-
-# In-memory dictionary to store OCR results
 OCR_RESULTS = {}
 
-# -------------------------------------------------------------------
-# HEALTH CHECK
-# -------------------------------------------------------------------
+
 @app.get("/health")
 async def health():
     return {"status": "ok"}
 
+
 # -------------------------------------------------------------------
-# UPLOAD IMAGE → GRIDFS
+# UPLOAD
 # -------------------------------------------------------------------
 @app.post("/upload")
 async def upload_image(file: UploadFile = File(...)):
     try:
         data = await file.read()
-        oid = await _gridfs.upload_from_stream(
+
+        fs = AsyncIOMotorGridFSBucket(_db, bucket_name=BUCKET_NAME)
+
+        oid = await fs.upload_from_stream(
             file.filename,
             data,
             metadata={"contentType": file.content_type},
@@ -54,70 +50,60 @@ async def upload_image(file: UploadFile = File(...)):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 # -------------------------------------------------------------------
-# HELPER: LOAD IMAGE FROM GRIDFS
+# LOAD IMAGE
 # -------------------------------------------------------------------
 async def load_image_from_gridfs(image_id: str):
     try:
         oid = ObjectId(image_id)
-        stream = await _gridfs.open_download_stream(oid)
+        fs = AsyncIOMotorGridFSBucket(_db, bucket_name=BUCKET_NAME)
+
+        stream = await fs.open_download_stream(oid)
         data = await stream.read()
         await stream.close()
         return data
-    except:
+    except Exception as e:
         raise HTTPException(status_code=404, detail="Image not found")
 
+
 # -------------------------------------------------------------------
-# GENERATE OCR FROM IMAGE ID
+# GENERATE OCR
 # -------------------------------------------------------------------
 @app.post("/generate/{image_id}")
 async def generate(image_id: str):
     raw_bytes = await load_image_from_gridfs(image_id)
 
-    # Convert to image
     img_array = np.frombuffer(raw_bytes, np.uint8)
     img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
 
     if img is None:
         raise HTTPException(status_code=400, detail="Unable to decode image")
 
-    # Run RapidOCR
     result, times = ocr_engine(img)
 
     if not result:
         raise HTTPException(status_code=500, detail="OCR returned empty result")
 
-    # Extract text
-    lines = []
-    for box, text, score in result:
-        lines.append(text)
-
-    full_text = "\n".join(lines)
+    extracted = [text for box, text, score in result]
+    full_text = "\n".join(extracted)
 
     OCR_RESULTS[image_id] = {
         "text": full_text,
         "details": result,
-        "timing": {
-            "det": times[0],
-            "cls": times[1],
-            "rec": times[2],
-        },
+        "timing": {"det": times[0], "cls": times[1], "rec": times[2]},
     }
 
     return {"image_id": image_id, "text": full_text}
 
-# -------------------------------------------------------------------
-# VIEW STORED OCR RESULT
-# -------------------------------------------------------------------
+
 @app.get("/view/{image_id}")
 async def view_details(image_id: str):
     if image_id not in OCR_RESULTS:
-        raise HTTPException(status_code=404, detail="No OCR result found. Use /generate first.")
+        raise HTTPException(status_code=404, detail="No OCR result found")
 
     return OCR_RESULTS[image_id]
 
-# -------------------------------------------------------------------
-# MAIN
-# -------------------------------------------------------------------
+
 if __name__ == "__main__":
     uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)