main.py - master - efr/SmartCity

main.py
194 строки7 KB
Sergey
add_project
27 фев 2025, 19:33
27 фев 2025, 19:337eeda08
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional, Dict, List, Any
import webbrowser
import uvicorn
import logging
import asyncio
import os
import subprocess

from utils.config_manager import load_tools, save_tools
from utils.yolo_service import yolo_processor, get_camera_state

import settings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory="static"), name="static")

def open_editor(file_path):
    if os.name == 'nt':
        subprocess.run(['notepad.exe', file_path])
    elif os.name == 'posix':
        if sys.platform == 'darwin':
            subprocess.run(['open', '-a', 'TextEdit', file_path])
        else:
            subprocess.run(['xdg-open', file_path])

class FrameRequest(BaseModel):
    frame: str
    camera_url: str

@app.post("/process")
async def process_frame(request: FrameRequest):
    try:
        result = await asyncio.to_thread(yolo_processor, request)
        return result
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

class ToolData(BaseModel):
    type: str
    coordinates: Dict[str, Any]
    class_ids: Optional[List[int]] = None
    class_id: Optional[int] = None
    camera_url: Optional[str] = None

class UndoRequest(BaseModel):
    camera_url: str

@app.post("/api/tools")
async def save_tool_data(data: ToolData):
    try:
        if not data.camera_url:
            raise ValueError("Camera URL is required")
        
        current_tools = load_tools(data.camera_url)
        
        if data.type == 'line':
            current_tools["lines"].append({
                "coordinates": data.coordinates,
                "class_ids": data.class_ids
            })
            current_tools["history"].append({
                "type": "line",
                "action": "add",
                "index": len(current_tools["lines"]) - 1
            })
        elif data.type == 'zone':
            current_tools["zones"].append({
                "coordinates": data.coordinates
            })
            current_tools["history"].append({
                "type": "zone",
                "action": "add",
                "index": len(current_tools["zones"]) - 1
            })
        elif data.type == 'traffic':
            prev_traffic = current_tools.get("traffic")
            prev_traffic_type = current_tools.get("traffic_type", 1)
            current_tools["traffic"] = {
                "coordinates": data.coordinates,
                "class_id": data.class_id
            }
            current_tools["traffic_type"] = data.class_id
            current_tools["history"].append({
                "type": "traffic",
                "action": "update",
                "prev_traffic": prev_traffic,
                "prev_traffic_type": prev_traffic_type
            })
        else:
            raise ValueError(f"Unknown tool type: {data.type}")
        
        save_tools(data.camera_url, current_tools)
        return JSONResponse({"status": "success", "data": data.dict()})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/tools/undo")
async def undo_last_action(data: UndoRequest):
    try:
        if not data.camera_url:
            raise ValueError("Camera URL is required")
        
        current_tools = load_tools(data.camera_url)
        history = current_tools.get("history", [])
        
        if not history:
            return JSONResponse({"status": "error", "message": "No actions to undo"}, status_code=400)
        
        last_action = history.pop()

        if last_action["type"] == "zone" and last_action["action"] == "add":
            index = last_action["index"]
            if 0 <= index < len(current_tools["zones"]):
                del current_tools["zones"][index]
            else:
                logger.warning(f"Zone index {index} is out of bounds.")
        elif last_action["type"] == "line" and last_action["action"] == "add":
            index = last_action["index"]
            if 0 <= index < len(current_tools["lines"]):
                del current_tools["lines"][index]
            else:
                logger.warning(f"Line index {index} is out of bounds.")
        elif last_action["type"] == "traffic" and last_action["action"] == "update":
            current_tools["traffic"] = last_action.get("prev_traffic")
            current_tools["traffic_type"] = last_action.get("prev_traffic_type", 1)
        else:
            logger.warning(f"Unsupported action for undo: {last_action}")
            return JSONResponse({"status": "error", "message": "Unsupported action"}, status_code=400)
        
        current_tools["history"] = history
        save_tools(data.camera_url, current_tools)
        return JSONResponse({"status": "success", "message": "Last action undone"})
    except Exception as e:
        logger.error(f"Failed to undo action: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

class StartRecordingRequest(BaseModel):
    camera_url: str

@app.post("/start-recording")
async def start_recording(request: StartRecordingRequest):
    try:
        camera_url = request.camera_url
        state = get_camera_state(camera_url)
        state['video_recorder'].start_recording(violation_type='vrec', width=1920, height=1080)
        return JSONResponse({"status": "success", "message": f"Recording started for {camera_url}"})
    except Exception as e:
        logger.error(f"Error starting recording: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))
    
@app.post("/settings")
async def read_settings():
    try:
        open_editor('settings.py')
        return JSONResponse({"status": "success"})
    except Exception as e:
        logger.error(f"Error starting recording: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))
    
@app.get("/")
async def read_index():
    return FileResponse("templates/index.html")

@app.get("/manifest.html")
async def read_manifest():
    return FileResponse("templates/manifest.html")

@app.get("/detection.html")
async def read_detection():
    return FileResponse("templates/detection.html")

@app.get("/target-classes")
async def get_target_classes():
    return sorted(settings.target_classes.items(), key=lambda x: x[0])

if __name__ == "__main__":
    webbrowser.open('http://localhost:8000')
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")