src/api.py

import os
import tempfile
from typing import List, Dict, Any

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse

from document_classifier import DocumentClassifier


app = FastAPI(title="Document Classifier API", version="1.0")

# Singleton classifier
classifier = DocumentClassifier()

# Expected document types provided by user
EXPECTED_LABELS: List[str] = [
    "CERTIFIED BIRTH CERTIFICATE",
    "CERTIFIED COURT DOCUMENT",
    "CERTIFIED HIGH SCHOOL TRANSCRIPT/REGISTRATION SIGNED BY SCHOOL OFFICIAL",
    "CERTIFIED MARRIAGE CERTIFICATE",
    "CERTIFIED NORTH CAROLINA MOTOR VEHICLE REPORT",
    "CERTIFIED OUT-OF-STATE MOTOR VEHICLE RECORD",
    "COMMERCIAL DRIVER LICENSE WITH PHOTO- VALID OR EXPIRED LESS THAN 1 YEAR",
    "DD2",
    "DD214",
    "DIVORCE DECREE",
    "IMMIGRATION DOCUMENT",
    "MILITARY DEPENDENT CARD",
    "NC DIPLOMA/GED FROM A COMMUNITY COLLEGE/COLLEGE/UNIVERSITY",
    "NC DL/ID WITH PHOTO VALID OR EXPIRED LESS THAN 2 YEARS",
    "NC LIMITED DRIVING PRIVILEGE ISSUED BY A COURT- VALID OR NOT EXPIRED MORE THAN ONE YEAR",
    "NC LP WITH PHOTO VALID OR EXPIRED LESS THAN 1 YEAR",
    "NON-CERTIFIED NORTH CAROLINA MOTOR VEHICLE REPORT",
    "OTHER",
    "OUT-OF-STATE DL/ID WITH PHOTO- VALID OR EXPIRED LESS THAN 2 YEARS",
    "OUT-OF-STATE LP WITH PHOTO- VALID OR EXPIRED LESS THAN 1 YEAR",
    "SOCIAL SECURITY CARD",
    "US VETERAN/ COMMON ACCESS CARD",
    "VALID MILITARY ID",
    "VALID PASSPORT",
    "ADDRESS AFFIDAVIT",
    "ANY DOCUMENT ISSUED BY THIS STATE OR COUNTY, CITY, OR THE FEDERAL GOVERNMENT",
    "LEASE OR HOUSING CONTRACT",
    "MILITARY ORDERS/DOCUMENTS",
    "MORTGAGE STATEMENT",
    "NC VEHICLE REGISTRATION CARD OR TITLE",
    "NC VOTER PRECINCT CARD",
    "OTHER",
    "PROPERTY OR INCOME  TAX STATEMENT",
    "UTILITY BILL OR CABLE BILL",
    "CERTIFICATE OF CITIZENSHIP",
    "CERTIFICATE OF NATURALIZATION",
    "DS2019 EXCHANGE VISITOR",
    "FOREIGN PASSPORT",
    "I-20 NONIMMIGRANT STUDENT CERTIFICATE",
    "I-220 B ORDER OF SUPERVISION W/ I-766 CARD",
    "I-512L AUTHORIZATION OF PAROLE OF AN ALIEN INTO U.S. W/SUPPORTING IMMIGRATION DOCUMENTS",
    "I-551 PERMANENT RESIDENT CARD",
    "I-766 EMPLOYMENT AUTHORIZATION CARD",
    "I-94 ARRIVAL/DEPARTURE RECORD",
    "MACHINE READABLE IMMIGRANT VISA WITH TEMPORARY I-551 LANGUAGE",
    "OTHER DOCUMENTS WITH ALIEN NUMBER",
    "TEMPORARY I-551 STAMP ON PASSPORT OR I-94",
    "UNEXPIRED FOREIGN PASSPORT W/CBP ADMISSION STAMP",
    "VALID US PASSPORT OR US PASSPORT CARD",
    "WAIVER TRAVELER/WAIVER BUSINESS (WT/WB) ADMISSION STAMP",
    "DL-123",
    "DL-123A-FLEET",
    "VALID VEHICLE INSURANCE BINDER",
    "VALID VEHICLE INSURANCE CARD",
    "VALID VEHICLE INSURANCE POLICY",
    "CERTIFICATE FOR SUPERVISING DRIVER (DL3B)",
    "DRIVER’S EDUCATION CERTIFICATE",
    "DRIVING ELIGIBILITY CERTIFICATE",
    "LEVEL 1 DRIVING LOG (DL4A)",
    "LEVEL 2 DRIVING LOG (DL4B)",
    "MINOR CONSENT FORM/ROAD TEST RESULTS (DL3A)",
    "OUT-OF-STATE TEEN DRIVER CONSENT FORM (DL3C)",
    "ROAD TEST EVALUATION FORM (DL4) (SKILLS TESTING)",
    "CDL 20",
    "CDL 5",
    "CDL MEDICAL CERTIFICATE",
    "CONDITION RESTORATION FORM",
    "COURT DOCUMENT – NAME CHANGE ORDER",
    "COURT DOCUMENT- GENDER CHANGE ORDER",
    "DL-77",
    "DL101",
    "DL300",
    "HEADGEAR AFFIDAVIT",
    "IGNITION INTERLOCK FORM",
    "MED CERT 1",
    "MEDICAL REPORT",
    "MOTORCYCLE ROAD TEST RESULTS (DL4)",
    "MSF CARD",
    "NSOPW FORM",
    "OTHER",
    "SKILLS TEST RESULTS",
    "1099 FORM",
    "DD-214",
    "PAYROLL RECORD/PAY STUB WITH FULL SOCIAL",
    "SOCIAL SECURITY CARD",
    "SOCIAL SECURITY DOCUMENT REFLECTING THE SSN",
    "VALID MILITARY RECORD WITH FULL SOCIAL",
    "W-2 FORM",
]


@app.post("/classify")
async def classify_document(file: UploadFile = File(...)) -> JSONResponse:
    if not file:
        raise HTTPException(status_code=400, detail="No file uploaded")

    suffix = os.path.splitext(file.filename or "uploaded")[1]
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            content = await file.read()
            tmp.write(content)
            tmp_path = tmp.name

        result = classifier.classify_document(tmp_path, allowed_labels=EXPECTED_LABELS)

        # Cleanup temp file
        try:
            os.unlink(tmp_path)
        except Exception:
            pass

        return JSONResponse(content={
            "success": result.get("success", False),
            "file_name": result.get("file_name"),
            "classification": result.get("classification"),
            "confidence": result.get("confidence"),
            "all_scores": result.get("all_scores"),
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))