Russian-ALPR / app /main.py
Naumeex's picture
Reliable employee seeding on startup
fa445c2
Raw
History Blame Contribute Delete
17.5 kB
"""FastAPI приложение для распознавания российских номеров."""
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
import cv2
from fastapi import FastAPI, UploadFile, File, Depends, HTTPException, Query
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from sqlalchemy import or_
from sqlalchemy.orm import Session
from fastapi import Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from app.core.config import (
DETECTOR_WEIGHTS, OCR_MODEL_NAME,
IMAGES_DIR, ANNOTATED_DIR, VIDEOS_DIR,
MAX_UPLOAD_SIZE_MB, ALLOWED_EXTENSIONS,
)
from app.db.database import init_db, get_session
from app.db.models import PlateRecord
from app.services.pipeline import PlatePipeline
from app.services.annotator import annotate_image
from app.db.models import PlateRecord, Employee
from app.services.access_control import find_employee_by_plate, make_access_decision
import tempfile
from app.services.video_processor import process_video, is_video_file
from app.core.config import ALLOWED_VIDEO_EXTENSIONS
# --- Lifespan: загрузка ML-моделей при старте ---
@asynccontextmanager
async def lifespan(app: FastAPI):
print("⏳ Инициализация БД...")
init_db()
# Автосидинг сотрудников при первом запуске (надёжный, без subprocess)
try:
from scripts.seed_employees import seed_if_empty
seed_if_empty()
except Exception as e:
print(f"⚠️ Автосидинг не выполнен: {e}")
print("⏳ Загрузка ML-моделей (~10 секунд)...")
app.state.pipeline = PlatePipeline(DETECTOR_WEIGHTS, ocr_model=OCR_MODEL_NAME)
print("✓ Готово к работе")
yield
print("👋 Завершаю работу")
app = FastAPI(
title="Russian ALPR",
description="Распознавание российских автомобильных номеров",
version="1.0.0",
lifespan=lifespan,
)
# Раздаём статику с картинками
app.mount("/images", StaticFiles(directory=str(IMAGES_DIR)), name="images")
app.mount("/annotated", StaticFiles(directory=str(ANNOTATED_DIR)), name="annotated")
templates = Jinja2Templates(directory="app/templates")
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/videos", StaticFiles(directory=str(VIDEOS_DIR)), name="videos")
# --- Эндпоинты ---
@app.get("/", response_class=HTMLResponse)
def page_home(request: Request):
return templates.TemplateResponse(
request=request,
name="index.html",
context={"active": "home"},
)
@app.get("/search", response_class=HTMLResponse)
def page_search(
request: Request,
q: str | None = None,
page: int = Query(1, ge=1),
session: Session = Depends(get_session),
):
per_page = 20
query = session.query(PlateRecord)
if q:
pattern = f"%{q.upper()}%"
query = query.filter(
or_(
PlateRecord.plate_text.like(pattern),
PlateRecord.plate_text_raw.like(pattern),
)
)
total = query.count()
plates = (
query.order_by(PlateRecord.detected_at.desc())
.offset((page - 1) * per_page)
.limit(per_page)
.all()
)
return templates.TemplateResponse(
request=request,
name="search.html",
context={
"active": "search",
"q": q,
"total": total,
"plates": [p.to_dict() for p in plates],
"page": page,
"per_page": per_page,
},
)
@app.get("/plates/{plate_id}", response_class=HTMLResponse)
def page_plate_detail(
plate_id: int,
request: Request,
session: Session = Depends(get_session),
):
record = session.query(PlateRecord).filter_by(id=plate_id).first()
if not record:
raise HTTPException(404, "Запись не найдена")
return templates.TemplateResponse(
request=request,
name="plate_detail.html",
context={
"active": None,
"plate": record.to_dict(),
},
)
@app.post("/api/upload")
async def upload_image(
file: UploadFile = File(...),
session: Session = Depends(get_session),
):
"""Загружает фото ИЛИ видео, распознаёт номер, сохраняет результат."""
ext = Path(file.filename or "").suffix.lower()
if ext not in ALLOWED_EXTENSIONS:
raise HTTPException(400, f"Неподдерживаемый формат: {ext}")
content = await file.read()
if len(content) > MAX_UPLOAD_SIZE_MB * 1024 * 1024:
raise HTTPException(413, f"Файл больше {MAX_UPLOAD_SIZE_MB} МБ")
is_video = ext in ALLOWED_VIDEO_EXTENSIONS
img_id = uuid.uuid4().hex
original_name = f"{img_id}.jpg" # всегда сохраняем итоговый кадр как jpg
annotated_name = f"{img_id}_annotated.jpg"
original_path = IMAGES_DIR / original_name
annotated_path = ANNOTATED_DIR / annotated_name
video_name = None # имя сохранённого видеофайла (для БД)
if is_video:
# --- Обработка видео ---
# Сохраняем видео постоянно (чтобы показывать в карточке)
video_name = f"{img_id}{ext}"
video_path = VIDEOS_DIR / video_name
video_path.write_bytes(content)
try:
result = process_video(app.state.pipeline, video_path,
frame_step=10, max_frames=60)
except ValueError as e:
video_path.unlink(missing_ok=True)
raise HTTPException(400, str(e))
if result is None:
video_path.unlink(missing_ok=True)
raise HTTPException(422, "В видео не найдено ни одного номера")
img = result["frame"]
detections = [result["detection"]]
cv2.imwrite(str(original_path), img)
else:
# --- Обработка фото ---
original_path.write_bytes(content)
img = cv2.imread(str(original_path))
if img is None:
original_path.unlink(missing_ok=True)
raise HTTPException(400, "Не удалось прочитать изображение")
detections = app.state.pipeline.process(img)
# Рисуем bbox
annotate_image(img, detections, annotated_path)
# Сохраняем в БД
saved_records = []
access_results = []
for d in detections:
record = PlateRecord(
image_filename=original_name,
annotated_filename=annotated_name,
plate_text=d["plate_text"],
plate_text_raw=d["plate_text_raw"],
is_valid_gost=d["is_valid_gost"],
bbox_x1=d["bbox"][0], bbox_y1=d["bbox"][1],
bbox_x2=d["bbox"][2], bbox_y2=d["bbox"][3],
confidence=d["confidence"],
is_fallback=d.get("fallback", False),
source_type="video" if is_video else "image",
video_filename=video_name,
)
session.add(record)
session.flush()
saved_records.append(record.to_dict())
emp = find_employee_by_plate(session, d["plate_text"])
access_results.append(make_access_decision(emp))
session.commit()
return {
"image_id": img_id,
"source_type": "video" if is_video else "image",
"original_url": f"/images/{original_name}",
"annotated_url": f"/annotated/{annotated_name}",
"detections": saved_records,
"access": access_results,
}
@app.get("/api/suggestions")
def get_suggestions(session: Session = Depends(get_session)):
"""Уникальные значения для автодополнения в форме сотрудника."""
departments = sorted({
d for (d,) in session.query(Employee.department).distinct().all() if d
})
positions = sorted({
p for (p,) in session.query(Employee.position).distinct().all() if p
})
names = sorted({
n for (n,) in session.query(Employee.full_name).distinct().all() if n
})
# Номера из реестра + из распознанных (история), чтобы можно было привязать существующий
reg_plates = {p for (p,) in session.query(Employee.plate_number).distinct().all() if p}
detected_plates = {
p for (p,) in session.query(PlateRecord.plate_text).distinct().all()
if p and len(p) >= 6
}
plates = sorted(reg_plates | detected_plates)
return {
"departments": departments,
"positions": positions,
"names": names,
"plates": plates,
}
@app.get("/api/plates")
def list_plates(
q: str | None = Query(None, description="Поиск по подстроке номера (case-insensitive)"),
limit: int = Query(50, ge=1, le=500),
offset: int = Query(0, ge=0),
session: Session = Depends(get_session),
):
"""Список записей с опциональным поиском по подстроке номера."""
query = session.query(PlateRecord)
if q:
# Ищем И по нормализованному (кириллица), И по сырому OCR (латиница)
# для случаев когда пользователь ввёл латиницей
pattern = f"%{q.upper()}%"
query = query.filter(
or_(
PlateRecord.plate_text.like(pattern),
PlateRecord.plate_text_raw.like(pattern),
)
)
total = query.count()
rows = (
query.order_by(PlateRecord.detected_at.desc())
.offset(offset)
.limit(limit)
.all()
)
return {
"total": total,
"limit": limit,
"offset": offset,
"query": q,
"results": [r.to_dict() for r in rows],
}
@app.get("/api/plates/{plate_id}")
def get_plate(plate_id: int, session: Session = Depends(get_session)):
"""Получить запись по ID."""
record = session.query(PlateRecord).filter_by(id=plate_id).first()
if not record:
raise HTTPException(404, f"Запись id={plate_id} не найдена")
return record.to_dict()
# --- Реестр сотрудников: API ---
@app.get("/api/employees")
def list_employees(
q: str | None = None,
department: str | None = None,
access: str | None = None, # "allowed" / "denied"
session: Session = Depends(get_session),
):
query = session.query(Employee)
if q:
pat = f"%{q.upper()}%"
query = query.filter(
or_(
Employee.full_name.like(f"%{q}%"),
Employee.plate_number.like(pat),
Employee.position.like(f"%{q}%"),
)
)
if department:
query = query.filter(Employee.department == department)
if access == "allowed":
query = query.filter(Employee.access_allowed == True) # noqa: E712
elif access == "denied":
query = query.filter(Employee.access_allowed == False) # noqa: E712
rows = query.order_by(Employee.full_name).all()
return {"total": len(rows), "results": [r.to_dict() for r in rows]}
@app.post("/api/employees")
def create_employee(data: dict, session: Session = Depends(get_session)):
required = {"full_name", "plate_number"}
if not required.issubset(data):
raise HTTPException(400, f"Обязательные поля: {required}")
emp = Employee(
full_name=data["full_name"].strip(),
position=data.get("position", "").strip() or None,
department=data.get("department", "").strip() or None,
plate_number=data["plate_number"].strip().upper(),
access_allowed=data.get("access_allowed", True),
notes=data.get("notes") or None,
)
session.add(emp)
session.commit()
session.refresh(emp)
return emp.to_dict()
@app.patch("/api/employees/{emp_id}")
def update_employee(emp_id: int, data: dict, session: Session = Depends(get_session)):
emp = session.query(Employee).filter_by(id=emp_id).first()
if not emp:
raise HTTPException(404, "Сотрудник не найден")
for field in ("full_name", "position", "department", "plate_number", "access_allowed", "notes"):
if field in data:
value = data[field]
if field == "plate_number" and isinstance(value, str):
value = value.strip().upper()
setattr(emp, field, value)
session.commit()
session.refresh(emp)
return emp.to_dict()
@app.delete("/api/employees/{emp_id}")
def delete_employee(emp_id: int, session: Session = Depends(get_session)):
emp = session.query(Employee).filter_by(id=emp_id).first()
if not emp:
raise HTTPException(404, "Сотрудник не найден")
session.delete(emp)
session.commit()
return {"deleted": emp_id}
@app.get("/api/check-access")
def check_access(
plate: str | None = None,
name: str | None = None,
position: str | None = None,
session: Session = Depends(get_session),
):
"""
Быстрая проверка: разрешён ли доступ человеку без фото.
Можно искать по номеру, ФИО или должности (любому одному из трёх).
Возвращает все совпадения с решением о доступе.
"""
if not any([plate, name, position]):
raise HTTPException(400, "Укажите хотя бы один параметр: plate, name или position")
query = session.query(Employee)
if plate:
# Нормализуем введённый номер (срежем регион если ввели) и ищем по подстроке
from app.services.plate_normalizer import normalize_plate
normalized, _ = normalize_plate(plate.strip())
search_value = normalized if normalized else plate.strip().upper()
query = query.filter(Employee.plate_number.like(f"%{search_value}%"))
if name:
query = query.filter(Employee.full_name.like(f"%{name.strip()}%"))
if position:
query = query.filter(Employee.position.like(f"%{position.strip()}%"))
rows = query.order_by(Employee.full_name).limit(20).all()
return {
"query": {"plate": plate, "name": name, "position": position},
"matches": [
{
"employee": e.to_dict(),
"access": make_access_decision(e),
}
for e in rows
],
}
# --- Реестр сотрудников: HTML ---
@app.get("/registry", response_class=HTMLResponse)
def page_registry(
request: Request,
q: str | None = None,
department: str | None = None,
access: str | None = None,
page: int = Query(1, ge=1),
session: Session = Depends(get_session),
):
per_page = 20
query = session.query(Employee)
if q:
pat = f"%{q.upper()}%"
query = query.filter(
or_(
Employee.full_name.like(f"%{q}%"),
Employee.plate_number.like(pat),
Employee.position.like(f"%{q}%"),
)
)
if department:
query = query.filter(Employee.department == department)
if access == "allowed":
query = query.filter(Employee.access_allowed == True) # noqa: E712
elif access == "denied":
query = query.filter(Employee.access_allowed == False) # noqa: E712
total = query.count()
employees = (
query.order_by(Employee.full_name)
.offset((page - 1) * per_page)
.limit(per_page)
.all()
)
departments = sorted({
d for (d,) in session.query(Employee.department).distinct().all() if d
})
return templates.TemplateResponse(
request=request,
name="registry.html",
context={
"active": "registry",
"employees": [e.to_dict() for e in employees],
"departments": departments,
"q": q or "",
"selected_dept": department or "",
"selected_access": access or "",
"total": total,
"page": page,
"per_page": per_page,
},
)
@app.delete("/api/plates/{plate_id}")
def delete_plate(plate_id: int, session: Session = Depends(get_session)):
"""Удалить запись и связанные файлы."""
record = session.query(PlateRecord).filter_by(id=plate_id).first()
if not record:
raise HTTPException(404, f"Запись id={plate_id} не найдена")
# Удаляем файлы (но осторожно: одна картинка может породить несколько записей)
same_image_count = session.query(PlateRecord).filter_by(
image_filename=record.image_filename
).count()
if same_image_count == 1: # это была последняя запись для этой картинки
(IMAGES_DIR / record.image_filename).unlink(missing_ok=True)
(ANNOTATED_DIR / record.annotated_filename).unlink(missing_ok=True)
session.delete(record)
session.commit()
return {"deleted": plate_id}