premdeep09/ANPR-System

1from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form
2from fastapi.middleware.cors import CORSMiddleware
3from pydantic import BaseModel
4from typing import List, Optional
5import os
6import shutil
7from database import DatabaseManager
8from pipeline import ANPRPipeline
9import logging
10
11logging.basicConfig(level=logging.INFO)
12logger = logging.getLogger("API")
13
14app = FastAPI(title="ANPR System API")
15
16# Add CORS Middleware
17app.add_middleware(
18    CORSMiddleware,
19    allow_origins=["*"],
20    allow_credentials=True,
21    allow_methods=["*"],
22    allow_headers=["*"],
23)
24
25# Initialize DB Manager
26db_manager = DatabaseManager(database="anpr_system")
27
28# Initialize AI Pipeline
29# We start this async/in a background thread inside the startup event
30# In production, RTSP URL can be read from environment variables
31RTSP_URL = os.getenv("RTSP_URL", "0") # Default to webcam for local tests
32pipeline = ANPRPipeline(
33    rtsp_url=RTSP_URL,
34    db_manager=db_manager,
35    storage_dir="storage",
36    cooldown_seconds=30 # 30 seconds cooldown to prevent rapid multi-detections of the same vehicle
37)
38
39@app.on_event("startup")
40def startup_event():
41    logger.info("Initializing ANPR Background Pipeline...")
42    pipeline.start()
43
44@app.on_event("shutdown")
45def shutdown_event():
46    logger.info("Shutting down ANPR Background Pipeline...")
47    pipeline.stop()
48
49# Data models
50class VehicleRecord(BaseModel):
51    id: int
52    plate_number: str
53    vehicle_type: Optional[str] = "Unknown"
54    entry_time: str
55    exit_time: Optional[str] = None
56    status: str
57    blacklisted: bool = False
58    plate_image_path: Optional[str] = None
59    vehicle_image_path: Optional[str] = None
60
61class ManualEntryRequest(BaseModel):
62    plate_number: str
63    vehicle_type: str
64
65from fastapi.responses import StreamingResponse
66import cv2
67import io
68
69@app.get("/")
70def read_root():
71    return {"message": "ANPR System API is running natively with YOLOv8 & MySQL!"}
72
73def generate_frames():
74    import time
75    while True:
76        frame = pipeline.latest_frame
77        if frame is None:
78            time.sleep(0.1)
79            continue
80        try:
81            ret, buffer = cv2.imencode('.jpg', frame)
82            if not ret: 
83                time.sleep(0.1)
84                continue
85            
86            yield (b'--frame\r\n'
87                   b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')
88            time.sleep(0.05)
89        except Exception as e:
90            time.sleep(0.1)
91
92@app.get("/api/video_feed")
93def video_feed():
94    return StreamingResponse(generate_frames(), media_type="multipart/x-mixed-replace; boundary=frame")
95
96@app.get("/api/latest_frame")
97def get_latest_frame():
98    frame = pipeline.latest_frame
99    if frame is None:
100        raise HTTPException(status_code=404, detail="No frame available")
101    _, buffer = cv2.imencode('.jpg', frame)
102    return StreamingResponse(io.BytesIO(buffer.tobytes()), media_type="image/jpeg")
103
104@app.get("/api/stats")
105def get_stats():
106    return db_manager.get_stats()
107
108@app.get("/api/vehicles", response_model=List[VehicleRecord])
109def get_vehicles(limit: int = 50):
110    return db_manager.get_recent_vehicles(limit=limit)
111
112@app.post("/api/manual_entry")
113def create_manual_entry(entry: ManualEntryRequest):
114    success = db_manager.add_or_update_vehicle(
115        plate_number=entry.plate_number.upper(),
116        vehicle_type=entry.vehicle_type
117    )
118    if success:
119        return {"message": f"Manual entry added for {entry.plate_number.upper()}"}
120    else:
121        raise HTTPException(status_code=500, detail="Failed to add manual entry to database.")
122
123@app.post("/api/video_source")
124async def change_video_source(
125    source_type: str = Form(...),
126    rtsp_url: Optional[str] = Form(None),
127    file: Optional[UploadFile] = File(None)
128):
129    if source_type == "Webcam":
130        new_source = "0"
131    elif source_type == "CCTV Stream":
132        if not rtsp_url:
133            raise HTTPException(status_code=400, detail="RTSP URL is required for CCTV Stream")
134        new_source = rtsp_url
135    elif source_type in ["Upload Video", "Upload Picture"]:
136        if not file:
137            raise HTTPException(status_code=400, detail="File is required for upload")
138        
139        upload_dir = os.path.join(os.getcwd(), "uploads")
140        os.makedirs(upload_dir, exist_ok=True)
141        file_path = os.path.join(upload_dir, file.filename)
142        
143        try:
144            content = await file.read()
145            with open(file_path, "wb") as buffer:
146                buffer.write(content)
147        except Exception as e:
148            raise HTTPException(status_code=500, detail=f"Failed to save file: {e}")
149            
150        new_source = file_path
151    else:
152        raise HTTPException(status_code=400, detail="Invalid source type")
153        
154    pipeline.set_source(new_source)
155    return {"message": f"Video source changed to: {source_type}"}
156
157if __name__ == "__main__":
158    import uvicorn
159    # Important: Setting reload=False is usually better when using background threads like YOLO
160    uvicorn.run(app, host="0.0.0.0", port=8000, reload=False)
161