Preload OCR model at startup

server.py

@@ -9,16 +9,18 @@ from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTa
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
-
-# Import existing functions
+from contextlib import asynccontextmanager
 from barcode_scanner import scan_barcode, scan_all_barcodes
-from ocr import read_chassis, postprocess_with_hint, ocr_image
+from ocr import read_chassis, postprocess_with_hint, ocr_image, get_ocr
 from preprocess import preprocess_chassis
 from evaluate import evaluate, get_pairs
 
-app = FastAPI(title="Chassis OCR API", description="API backend for Chassis OCR PWA")
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    get_ocr()
+    yield
 
-# CORS middleware for testing
+app = FastAPI(title="Chassis OCR API", description="API backend for Chassis OCR PWA", lifespan=lifespan)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -33,7 +35,6 @@ CONFIG_PATH = "config.json"
 os.makedirs(TEMP_DIR, exist_ok=True)
 os.makedirs(RESULTS_DIR, exist_ok=True)
 
-# Helper to convert cv2 image to base64 jpeg
 def cv2_to_base64(img):
     _, buffer = cv2.imencode('.jpg', img)
     return base64.b64encode(buffer).decode('utf-8')
@@ -51,7 +52,6 @@ def get_test_pairs():
     chassis_dir = "images/chassis"
     if not os.path.exists(barcode_dir) or not os.path.exists(chassis_dir):
         return []
-    
     barcodes = {os.path.splitext(f)[0] for f in os.listdir(barcode_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))}
     chassis = {os.path.splitext(f)[0] for f in os.listdir(chassis_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))}
     common = sorted(list(barcodes & chassis))
@@ -59,7 +59,6 @@ def get_test_pairs():
 
 @app.get("/api/scan-barcode/{key}")
 def scan_barcode_by_key(key: str):
-    """Scan a barcode image from the dataset by its key (filename without extension)."""
     for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.PNG']:
         path = os.path.join("images/barcode", f"{key}{ext}")
         if os.path.exists(path):
@@ -67,7 +66,6 @@ def scan_barcode_by_key(key: str):
             return {"success": result is not None, "barcode": result}
     raise HTTPException(status_code=404, detail=f"Barcode image for key '{key}' not found")
 
-
 @app.post("/api/scan-barcode")
 async def api_scan_barcode(file: UploadFile = File(...)):
     temp_filename = f"{uuid.uuid4()}_{file.filename}"
@@ -75,7 +73,6 @@ async def api_scan_barcode(file: UploadFile = File(...)):
     try:
         with open(temp_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
-        
         result = scan_barcode(temp_path)
         return {"success": result is not None, "barcode": result}
     except Exception as e:
@@ -100,7 +97,6 @@ async def save_config(config_data: dict):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# Global state to keep track of batch evaluation runs
 eval_status = {"running": False, "progress": 0, "total": 0, "results": []}
 
 def run_evaluation_task():
@@ -108,10 +104,7 @@ def run_evaluation_task():
     try:
         eval_status["running"] = True
         eval_status["progress"] = 0
-        
-        # We can call the evaluate function but let's read the report afterwards
         evaluate()
-        
         report_path = os.path.join(RESULTS_DIR, "report.json")
         if os.path.exists(report_path):
             with open(report_path) as f:
@@ -128,7 +121,6 @@ def trigger_evaluation(background_tasks: BackgroundTasks):
     global eval_status
     if eval_status["running"]:
         return {"status": "already_running", "message": "Evaluation task is currently running"}
-    
     eval_status = {"running": True, "progress": 0, "total": 50, "results": []}
     background_tasks.add_task(run_evaluation_task)
     return {"status": "started", "message": "Batch evaluation started in the background"}
@@ -143,7 +135,6 @@ def get_evaluation_status():
                 results = json.load(f)
         except Exception:
             pass
-            
     return {
         "running": eval_status["running"],
         "progress": eval_status["progress"],
@@ -165,8 +156,6 @@ async def match_chassis(
     temp_path = None
 
     if chassis_key:
-        # Load from test set
-        # Check standard extensions (.jpg, .png, etc.)
         for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.PNG']:
             p = os.path.join("images/chassis", f"{chassis_key}{ext}")
             if os.path.exists(p):
@@ -175,7 +164,6 @@ async def match_chassis(
         if not chassis_path:
             raise HTTPException(status_code=404, detail=f"Chassis image for key '{chassis_key}' not found in images/chassis")
     else:
-        # Save uploaded file
         temp_filename = f"{uuid.uuid4()}_{chassis_file.filename}"
         temp_path = os.path.join(TEMP_DIR, temp_filename)
         with open(temp_path, "wb") as buffer:
@@ -183,18 +171,12 @@ async def match_chassis(
         chassis_path = temp_path
 
     try:
-        # 1. Run Preprocessing to get variations
-        # Use save_comparison=True so we also write the side-by-side view to results/ (useful for viewing static file later)
-        # Note: if it's a temp file, let's create a friendly name for results comparison
         save_comp = True
         comp_filename = chassis_key if chassis_key else os.path.splitext(chassis_file.filename)[0]
-        
         variations = preprocess_chassis(chassis_path, save_comparison=save_comp)
-        
-        # 2. Get base64 representation of original and each variation
         original_img = cv2.imread(chassis_path)
         base64_original = cv2_to_base64(original_img)
-        
+
         base64_variations = []
         labels = ["CLAHE", "Bilateral", "Otsu", "Adaptive"]
         for idx, var in enumerate(variations):
@@ -203,11 +185,10 @@ async def match_chassis(
                 "base64": cv2_to_base64(var)
             })
 
-        # 3. Run OCR on each variation and compute scores, finding the best
         best_text, best_conf, best_score = "", 0.0, -1
         winning_label = ""
         variation_details = []
-        
+
         for idx, var in enumerate(variations):
             text, conf = ocr_image(var)
             score = conf * max(len(text), 1)
@@ -221,16 +202,14 @@ async def match_chassis(
                 best_text, best_conf, best_score = text, conf, score
                 winning_label = labels[idx]
 
-        # 4. Perform error correction and match check
         corrected_text, is_match = postprocess_with_hint(best_text, barcode_val)
-        
         status = "FAILED"
         if best_text == barcode_val:
             status = "EXACT"
         elif is_match:
             status = "CORRECTED"
 
-        response_data = {
+        return {
             "success": is_match,
             "status": status,
             "barcode_val": barcode_val,
@@ -243,7 +222,6 @@ async def match_chassis(
             "variation_details": variation_details,
             "comparison_url": f"/results/{comp_filename}_comparison.jpg" if save_comp else None
         }
-        return response_data
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -251,14 +229,9 @@ async def match_chassis(
         if temp_path and os.path.exists(temp_path):
             os.remove(temp_path)
 
-# Serve results images directly
 app.mount("/results", StaticFiles(directory="results"), name="results")
-
-# Serve frontend application static files
-# We will mount at "/" with html=True so index.html is served automatically
-# Make sure to run this *after* route declarations
 app.mount("/", StaticFiles(directory="web", html=True), name="static")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=True)
+    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=True)