server.py

import os
import json
import shutil
import base64
import uuid
import cv2
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from barcode_scanner import scan_barcode, scan_all_barcodes
from ocr import read_chassis, postprocess_with_hint, ocr_image, get_ocr
from preprocess import preprocess_chassis
from evaluate import evaluate, get_pairs

@asynccontextmanager
async def lifespan(app: FastAPI):
    get_ocr()
    yield

app = FastAPI(title="Chassis OCR API", description="API backend for Chassis OCR PWA", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

TEMP_DIR = "temp_uploads"
RESULTS_DIR = "results"
CONFIG_PATH = "config.json"
os.makedirs(TEMP_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)

def cv2_to_base64(img):
    _, buffer = cv2.imencode('.jpg', img)
    return base64.b64encode(buffer).decode('utf-8')

@app.get("/api/status")
def get_status():
    return {
        "status": "online",
        "message": "OCR Backend is active"
    }

@app.get("/api/test-pairs")
def get_test_pairs():
    barcode_dir = "images/barcode"
    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))
    return common

@app.get("/api/scan-barcode/{key}")
def scan_barcode_by_key(key: str):
    for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.PNG']:
        path = os.path.join("images/barcode", f"{key}{ext}")
        if os.path.exists(path):
            result = scan_barcode(path)
            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}"
    temp_path = os.path.join(TEMP_DIR, temp_filename)
    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:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

@app.get("/api/config")
def get_config():
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH) as f:
            return json.load(f)
    return {}

@app.post("/api/config")
async def save_config(config_data: dict):
    try:
        with open(CONFIG_PATH, "w") as f:
            json.dump(config_data, f, indent=2)
        return {"status": "success", "message": "Configuration updated successfully"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

eval_status = {"running": False, "progress": 0, "total": 0, "results": []}

def run_evaluation_task():
    global eval_status
    try:
        eval_status["running"] = True
        eval_status["progress"] = 0
        evaluate()
        report_path = os.path.join(RESULTS_DIR, "report.json")
        if os.path.exists(report_path):
            with open(report_path) as f:
                eval_status["results"] = json.load(f)
        eval_status["progress"] = len(eval_status["results"])
        eval_status["total"] = len(eval_status["results"])
    except Exception as e:
        print(f"Error in evaluation background task: {e}")
    finally:
        eval_status["running"] = False

@app.post("/api/evaluate")
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"}

@app.get("/api/evaluate/status")
def get_evaluation_status():
    report_path = os.path.join(RESULTS_DIR, "report.json")
    results = []
    if os.path.exists(report_path):
        try:
            with open(report_path) as f:
                results = json.load(f)
        except Exception:
            pass
    return {
        "running": eval_status["running"],
        "progress": eval_status["progress"],
        "total": eval_status["total"],
        "has_existing_report": len(results) > 0,
        "results": results if not eval_status["running"] else eval_status["results"]
    }

@app.post("/api/match")
async def match_chassis(
    barcode_val: str = Form(...),
    chassis_file: UploadFile = File(None),
    chassis_key: str = Form(None)
):
    if not chassis_file and not chassis_key:
        raise HTTPException(status_code=400, detail="Either chassis_file or chassis_key must be provided")

    chassis_path = None
    temp_path = None

    if chassis_key:
        for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.PNG']:
            p = os.path.join("images/chassis", f"{chassis_key}{ext}")
            if os.path.exists(p):
                chassis_path = p
                break
        if not chassis_path:
            raise HTTPException(status_code=404, detail=f"Chassis image for key '{chassis_key}' not found in images/chassis")
    else:
        temp_filename = f"{uuid.uuid4()}_{chassis_file.filename}"
        temp_path = os.path.join(TEMP_DIR, temp_filename)
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(chassis_file.file, buffer)
        chassis_path = temp_path

    try:
        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)
        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):
            base64_variations.append({
                "label": labels[idx],
                "base64": cv2_to_base64(var)
            })

        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)
            variation_details.append({
                "label": labels[idx],
                "text": text,
                "confidence": conf,
                "score": score
            })
            if score > best_score:
                best_text, best_conf, best_score = text, conf, score
                winning_label = labels[idx]

        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"

        return {
            "success": is_match,
            "status": status,
            "barcode_val": barcode_val,
            "raw_ocr": best_text,
            "corrected_ocr": corrected_text,
            "confidence": best_conf,
            "winning_label": winning_label,
            "variations": base64_variations,
            "original": base64_original,
            "variation_details": variation_details,
            "comparison_url": f"/results/{comp_filename}_comparison.jpg" if save_comp else None
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)

app.mount("/results", StaticFiles(directory="results"), name="results")
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)