app.py

import numpy as np
import torch
import torchvision
import my_train
from camera import ThreadedCamera

import time
import threading
from flask import (
    Flask,
    jsonify,
    render_template,
    send_from_directory,
    request,
    Response,
)
from flask_socketio import SocketIO
from flask_sqlalchemy import SQLAlchemy
import cv2
from ultralytics import YOLO
import json
import os
from datetime import datetime

app = Flask(__name__, static_folder="static")
socketio = SocketIO(app, cors_allowed_origins="*")

# 调试模式
debug = True


# 添加 WebSocket 事件处理
@socketio.on("connect")
def handle_connect():
    print("客户端已连接")


@socketio.on("disconnect")
def handle_disconnect():
    print("客户端已断开连接")


# 数据库配置
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///cv_stats.db"
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
db = SQLAlchemy(app)


# 统计数据模型
class CharacterStats(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    character_name = db.Column(db.String(100), nullable=False)
    count = db.Column(db.Integer, default=0)
    timestamp = db.Column(db.DateTime, default=datetime.now)


class RegionStats(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    region_name = db.Column(db.String(100), nullable=False)
    count = db.Column(db.Integer, default=0)
    timestamp = db.Column(db.DateTime, default=datetime.now)


# 全局变量
camera = None
is_processing = False
inference_thread = None
stop_event = threading.Event()

# 创建保存图片的目录
SAVE_DIR = "saved_images"
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)


@app.route("/")
def index():
    """返回前端页面"""
    return render_template("index.html")  # 渲染 templates 目录下的 HTML 模板


@app.route("/api/startProcessing", methods=["POST"])
def start_processing():
    """启动模型推理线程"""
    global is_processing, inference_thread, stop_event
    if not is_processing:
        print("已启动推理线程")
        is_processing = True
        stop_event.clear()
        inference_thread = threading.Thread(target=run_inference_loop)
        inference_thread.start()
    return jsonify({"status": "processing started"}), 200


@app.route("/api/stopProcessing", methods=["POST"])
def stop_processing():
    """停止模型推理线程"""
    global is_processing, inference_thread, stop_event
    if is_processing:
        is_processing = False
        stop_event.set()  # 通知线程退出
        if inference_thread is not None:
            inference_thread.join()
            inference_thread = None
    return jsonify({"status": "processing stopped"}), 200


@app.route("/api/camera_info")
def get_camera_info():
    """获取摄像头信息"""
    global camera
    if camera:
        return jsonify({"width": camera.frame_width, "height": camera.frame_height})
    return jsonify({"error": "Camera not initialized"}), 404


@app.route("/api/name_mapping")
def get_name_mapping():
    """获取英文名到中文名的映射关系"""
    try:
        with open("assets/labels.json", "r", encoding="utf-8") as f:
            name_mapping = json.load(f)
        return jsonify(name_mapping), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


def gen_frames():
    """生成视频流"""
    global camera
    while True:
        frame = camera.read()
        ret, buffer = cv2.imencode(".jpg", frame)
        frame_bytes = buffer.tobytes()
        yield (
            b"--frame\r\n" b"Content-Type: image/jpeg\r\n\r\n" + frame_bytes + b"\r\n"
        )


@app.route("/video_feed")
def video_feed():
    return Response(gen_frames(), mimetype="multipart/x-mixed-replace; boundary=frame")


def load_model(method="cifar10"):
    """加载模型"""
    # if method == "cifar10":
    model_path = "model/mytrain-cifar10.pth"
    model = my_train.load_full_model(model_path).to(my_train.device)
    model.eval()
    with open("assets/labels.json", "r", encoding="utf-8") as f:
        labels = json.load(f)
    return model, labels, my_train.device
    # else:
    #     return YOLO("model/yolo11m-02-01-best.pt", task="classify", verbose=False)


def center_crop_resize(image, ratio, size):
    """将图像按比例中心裁剪并调整大小"""
    h, w, _ = image.shape
    crop_size = int(min(h, w) * ratio)
    center_h, center_w = h // 2, w // 2
    start_h, start_w = center_h - crop_size // 2, center_w - crop_size // 2
    cropped = image[start_h : start_h + crop_size, start_w : start_w + crop_size]
    resized = cv2.resize(cropped, size)
    return resized


def run_inference_loop():
    """后台线程：进行模型推理"""
    global camera
    try:
        model, lables, device = load_model(method="cifar10")
        # 定义预处理转换
        transform = torchvision.transforms.Compose(
            [
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize(
                    [0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]
                ),  # from ImageNet 微调
            ]
        )

        while not stop_event.is_set():
            frame = camera.read()
            frame = center_crop_resize(frame, 0.7, (32, 32))
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            input_tensor = transform(frame_rgb).half().unsqueeze(0).to(device)
            # 进行模型推理
            with torch.no_grad():
                output = model(input_tensor)
            probabilities = torch.softmax(output, dim=1)[0].cpu().numpy()
            top5_indices = np.argsort(probabilities)[-5:][::-1]
            top5_conf = probabilities[top5_indices].tolist()

            results = [
                {
                    "className": str(idx),
                    "classNameCN": lables[str(idx)],
                    "probability": round(conf, 4),
                }
                for idx, conf in zip(top5_indices, top5_conf)
            ]
            if debug:
                print(f"推理：{lables[str(top5_indices[0])]}")

            socketio.emit("inference_result", {"top5": results})
            time.sleep(0.8)  # 控制推理频率

    except Exception as e:
        print(f"推理循环错误: {type(e).__name__} - {e}")
        import traceback

        traceback.print_exc()


@app.route("/api/save_image", methods=["POST"])
def save_image():
    """保存图片到本地"""
    try:
        # 获取Base64编码的图片数据
        data = request.get_json()
        image_data = data.get("image")
        if not image_data:
            return jsonify({"error": "No image data provided"}), 400

        # 解码Base64图片数据
        import base64

        image_data = base64.b64decode(image_data.split(",")[1])

        # 生成文件名（使用当前时间）
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}.png"
        filepath = os.path.join(SAVE_DIR, filename)
        with open(filepath, "wb") as f:
            f.write(image_data)

        # 通过WebSocket发送图片保存通知
        socketio.emit("image_saved", {"filename": filename, "timestamp": timestamp})

        return (
            jsonify({"status": "success", "filename": filename, "filepath": filepath}),
            200,
        )

    except Exception as e:
        return jsonify({"error": str(e)}), 500


# 添加静态路由
@app.route("/saved_images/<path:filename>")
def serve_saved_image(filename):
    """提供保存的截图"""
    return send_from_directory(SAVE_DIR, filename)


@app.route("/api/update_stats", methods=["POST"])
def update_stats():
    """更新统计数据"""
    try:
        data = request.get_json()
        character = data.get("character")
        region = data.get("region")

        # 更新角色统计
        if character:
            char_stat = CharacterStats.query.filter_by(character_name=character).first()
            if char_stat:
                char_stat.count += 1
            else:
                char_stat = CharacterStats(character_name=character, count=1)
                db.session.add(char_stat)

        # 更新地域统计
        if region:
            region_stat = RegionStats.query.filter_by(region_name=region).first()
            if region_stat:
                region_stat.count += 1
            else:
                region_stat = RegionStats(region_name=region, count=1)
                db.session.add(region_stat)

        db.session.commit()
        return jsonify({"status": "success"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/api/get_stats")
def get_stats():
    """获取统计数据"""
    try:
        # 获取角色统计
        char_stats = (
            CharacterStats.query.order_by(CharacterStats.count.desc()).limit(5).all()
        )
        char_stats_data = [
            {"name": stat.character_name, "count": stat.count} for stat in char_stats
        ]

        # 获取地域统计
        region_stats = (
            RegionStats.query.order_by(RegionStats.count.desc()).limit(5).all()
        )
        region_stats_data = [
            {"name": stat.region_name, "count": stat.count} for stat in region_stats
        ]

        return (
            jsonify({"characters": char_stats_data, "regions": region_stats_data}),
            200,
        )
    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    try:
        # 创建数据库表
        with app.app_context():
            db.create_all()

        if debug:
            source = 0
        else:
            print("请选择摄像头源：")
            print("0: 使用电脑自带摄像头")
            print("1: 使用外接摄像头")
            while True:
                try:
                    source = int(input("请输入摄像头源(0 或 1):"))
                    if source == 0 or source == 1:
                        break
                    else:
                        print("无效的选择，请输入 0 或 1")
                except ValueError:
                    print("请输入有效的数字")
        print("Starting camera...")
        camera = ThreadedCamera(camera_id=source)
        camera.start()

        print(f"Load name converting...")
        global cvt_name_to_cn
        with open("assets/labels.json", "r", encoding="utf-8") as f:
            cvt_name_to_cn = json.load(f)

        socketio.run(app, debug=True, host="0.0.0.0", port=2050, use_reloader=False)
    finally:
        if camera:
            camera.release()