diff --git a/data/training/README.dataset.txt b/data/training/README.dataset.txt new file mode 100644 index 0000000..28b437f --- /dev/null +++ b/data/training/README.dataset.txt @@ -0,0 +1,6 @@ +# DNF > 2023-06-06 3:57pm +https://universe.roboflow.com/atomade/dnf-fzffw + +Provided by a Roboflow user +License: CC BY 4.0 + diff --git a/data/training/README.roboflow.txt b/data/training/README.roboflow.txt new file mode 100644 index 0000000..0828d7c --- /dev/null +++ b/data/training/README.roboflow.txt @@ -0,0 +1,36 @@ + +DNF - v1 2023-06-06 3:57pm +============================== + +This dataset was exported via roboflow.com on June 6, 2023 at 8:01 AM GMT + +Roboflow is an end-to-end computer vision platform that helps you +* collaborate with your team on computer vision projects +* collect & organize images +* understand and search unstructured image data +* annotate, and create datasets +* export, train, and deploy computer vision models +* use active learning to improve your dataset over time + +For state of the art Computer Vision training notebooks you can use with this dataset, +visit https://github.com/roboflow/notebooks + +To find over 100k other datasets and pre-trained models, visit https://universe.roboflow.com + +The dataset includes 649 images. +Player-gold-monster-boss are annotated in YOLOv8 format. + +The following pre-processing was applied to each image: +* Auto-orientation of pixel data (with EXIF-orientation stripping) +* Resize to 640x640 (Stretch) + +The following augmentation was applied to create 3 versions of each source image: +* Randomly crop between 0 and 40 percent of the image +* Random brigthness adjustment of between -30 and +30 percent +* Random exposure adjustment of between -20 and +20 percent +* Random Gaussian blur of between 0 and 1 pixels + +The following transformations were applied to the bounding boxes of each image: +* Salt and pepper noise was applied to 5 percent of pixels + + diff --git a/data/training/data.yaml b/data/training/data.yaml index 843b7d9..e3c6ede 100644 --- a/data/training/data.yaml +++ b/data/training/data.yaml @@ -1,16 +1,12 @@ path: /workspace/dnf-auto-cloud/data/training -train: /workspace/dnf-auto-cloud/data/training/images/train -val: /workspace/dnf-auto-cloud/data/training/images/val -test: /workspace/dnf-auto-cloud/data/training/images/test -nc: 10 -names: -- monster -- boss -- door -- item -- npc -- player -- hp_bar -- mp_bar -- skill_ready -- cooldown +train: images/train +val: images/val +test: images/test +nc: 15 +names: ['door_boss', 'door_boss_activated', 'door_normal', 'door_normal_activated', 'item_diamondcoin_advanced', 'item_gold', 'monster_hebron_archer', 'monster_hebron_boss', 'monster_hebron_bulbhead', 'monster_hebron_robot', 'monster_hebron_stone', 'monster_hebron_tortoise', 'player_spectre', 'window_continue', 'window_result'] +roboflow: + workspace: atomade + project: dnf-fzffw + version: 1 + license: CC BY 4.0 + url: https://universe.roboflow.com/atomade/dnf-fzffw/dataset/1 \ No newline at end of file diff --git a/dnf-client-package/README.txt b/dnf-client-package/README.txt deleted file mode 100644 index d8079d5..0000000 --- a/dnf-client-package/README.txt +++ /dev/null @@ -1,9 +0,0 @@ -DNF自动化客户端 - -安装说明: -1. 安装Python 3.8或更高版本 -2. 安装依赖包: - pip install pillow numpy websockets keyboard mouse pywin32 mss - -3. 编辑config.ini设置服务器地址 -4. 运行start.bat启动客户端 diff --git a/dnf-client-package/client.py b/dnf-client-package/client.py deleted file mode 100644 index 68aec10..0000000 --- a/dnf-client-package/client.py +++ /dev/null @@ -1,837 +0,0 @@ -#!/usr/bin/env python -# -*- coding: utf-8 -*- - -""" -DNF自动化客户端，负责截取游戏画面并执行服务器返回的操作 -优化版 - 增加断线重连、性能优化和增强的游戏状态管理 -""" - -import os -import sys -import json -import base64 -import time -import random -import asyncio -import websockets -import configparser -import ssl -import ctypes -from io import BytesIO -from datetime import datetime -import threading -import logging -import traceback - -# 图像处理 -from PIL import Image -import numpy as np - -# Windows API -import win32gui -import win32con -import win32api -import win32process - -# 输入模拟 -import keyboard -import mouse - -# 尝试导入mss库，如果不存在则继续使用PIL -try: - import mss - MSS_AVAILABLE = True -except ImportError: - MSS_AVAILABLE = False - print("警告: mss库未安装，将使用PIL进行截图（性能较低）") - print("请使用 pip install mss 安装以获得更好的性能") - -# 配置文件路径 - 使用绝对路径 -CONFIG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.ini") - -# 日志设置 -logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=[ - logging.FileHandler(os.path.join(os.path.dirname(os.path.abspath(__file__)), "client.log"), encoding="utf-8"), - logging.StreamHandler() - ] -) -logger = logging.getLogger("DNFAutoClient") - -class DNFAutoClient: - """DNF自动化客户端类""" - - def __init__(self): - """初始化客户端""" - self.config = self.load_config() - self.server_url = self.config.get("Server", "url") - self.client_id = None - self.running = False - self.ws = None - self.capture_interval = float(self.config.get("Capture", "interval")) - self.max_retries = int(self.config.get("Connection", "max_retries", fallback="5")) - self.retry_delay = int(self.config.get("Connection", "retry_delay", fallback="5")) - - # 增强的游戏状态 - self.game_state = { - "in_battle": False, - "current_map": "", - "hp_percent": 100, - "mp_percent": 100, - "active_buffs": [], - "cooldowns": {}, - "inventory_full": False, - "current_quest": None, - "last_operation_time": time.time(), - "session_start_time": time.time() - } - - # 连接状态 - self.last_heartbeat_time = 0 - self.connection_attempts = 0 - self.reconnecting = False - - def load_config(self): - """加载配置文件""" - if not os.path.exists(CONFIG_FILE): - self.create_default_config() - - config = configparser.ConfigParser() - config.read(CONFIG_FILE, encoding="utf-8") - return config - - def create_default_config(self): - """创建默认配置文件""" - config = configparser.ConfigParser() - - config["Server"] = { - "url": "wss://your-server-url:8080/ws", - "verify_ssl": "false" - } - - config["Capture"] = { - "interval": "0.5", - "quality": "70" - } - - config["Game"] = { - "window_title": "地下城与勇士", - "key_mapping": "default" - } - - config["Connection"] = { - "max_retries": "5", - "retry_delay": "5", - "heartbeat_interval": "5" - } - - # 保存配置 - with open(CONFIG_FILE, "w", encoding="utf-8") as f: - config.write(f) - - logger.info(f"已创建默认配置文件: {CONFIG_FILE}") - - async def connect(self): - """连接到服务器""" - logger.info(f"正在连接到服务器: {self.server_url}") - - ssl_context = None - if self.server_url.startswith("wss://"): - ssl_context = ssl.create_default_context() - if self.config.get("Server", "verify_ssl").lower() == "false": - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE - - try: - self.ws = await websockets.connect( - self.server_url, - ssl=ssl_context, - max_size=10 * 1024 * 1024, # 10MB - ping_interval=None, # 禁用自动ping，我们将使用自己的心跳 - close_timeout=5 - ) - - # 等待认证 - await self.authenticate() - - logger.info("已连接到服务器") - self.connection_attempts = 0 # 重置连接尝试次数 - self.last_heartbeat_time = time.time() - return True - - except Exception as e: - logger.error(f"连接服务器失败: {e}") - return False - - async def connect_with_retry(self): - """带重试机制的连接函数""" - if self.reconnecting: - logger.info("已有重连过程在进行中，跳过") - return False - - self.reconnecting = True - retry_count = 0 - - try: - while retry_count < self.max_retries and not self.ws: - try: - logger.info(f"尝试连接服务器 (尝试 {retry_count + 1}/{self.max_retries})...") - success = await self.connect() - if success: - self.reconnecting = False - return True - except Exception as e: - logger.error(f"连接服务器失败: {e}") - - retry_count += 1 - retry_delay = min(60, self.retry_delay * (2 ** retry_count)) # 指数退避策略 - logger.info(f"等待 {retry_delay} 秒后重试...") - await asyncio.sleep(retry_delay) - - if not self.ws: - logger.error(f"达到最大重试次数 ({self.max_retries})，连接失败") - self.reconnecting = False - return False - except Exception as e: - logger.error(f"重连过程中发生错误: {e}") - self.reconnecting = False - return False - - self.reconnecting = False - return True - - async def authenticate(self): - """客户端认证""" - try: - # 等待认证挑战 - challenge_raw = await self.ws.recv() - challenge = json.loads(challenge_raw) - - if challenge.get("type") != "auth_challenge": - raise ValueError("无效的认证挑战") - - # 收集系统信息 - system_info = self.get_system_info() - - # 发送认证响应 - await self.ws.send(json.dumps({ - "type": "auth_response", - "response": f"client_{random.getrandbits(32)}", - "client_info": { - "version": "1.0.1", # 更新版本号 - "os": "Windows", - "screen_resolution": self.get_screen_resolution(), - "system_info": system_info - } - })) - - # 等待认证结果 - result_raw = await self.ws.recv() - result = json.loads(result_raw) - - if result.get("type") != "auth_result" or result.get("status") != "success": - raise ValueError("认证失败") - - # 保存客户端ID - self.client_id = result.get("client_id") - logger.info(f"认证成功，客户端ID: {self.client_id}") - - except Exception as e: - logger.error(f"认证失败: {e}") - raise - - def get_system_info(self): - """获取系统信息""" - system_info = {} - try: - system_info["hostname"] = os.environ.get("COMPUTERNAME", "Unknown") - system_info["username"] = os.environ.get("USERNAME", "Unknown") - system_info["processor"] = os.environ.get("PROCESSOR_IDENTIFIER", "Unknown") - - # 获取系统内存信息 - mem = ctypes.c_ulonglong() - ctypes.windll.kernel32.GetPhysicallyInstalledSystemMemory(ctypes.byref(mem)) - system_info["memory_gb"] = round(mem.value / (1024 * 1024), 2) - - except Exception as e: - logger.error(f"获取系统信息失败: {e}") - - return system_info - - def get_screen_resolution(self): - """获取屏幕分辨率""" - user32 = ctypes.windll.user32 - return [user32.GetSystemMetrics(0), user32.GetSystemMetrics(1)] - - def get_game_window(self): - """获取游戏窗口句柄""" - window_title = self.config.get("Game", "window_title") - hwnd = win32gui.FindWindow(None, window_title) - if hwnd == 0: - logger.warning(f"找不到游戏窗口: {window_title}") - return None - return hwnd - - def capture_game_screen(self, hwnd=None): - """截取游戏画面""" - try: - if hwnd is None: - hwnd = self.get_game_window() - if hwnd is None: - return None - - # 获取窗口位置和大小 - rect = win32gui.GetWindowRect(hwnd) - x, y, width, height = rect[0], rect[1], rect[2] - rect[0], rect[3] - rect[1] - - # 检查窗口是否最小化 - if width <= 0 or height <= 0: - logger.warning("游戏窗口被最小化或大小无效") - return None - - # 使用mss进行截图(如果可用)，速度比PIL快5-10倍 - if MSS_AVAILABLE: - with mss.mss() as sct: - monitor = {"top": y, "left": x, "width": width, "height": height} - sct_img = sct.grab(monitor) - # 将mss图像转换为PIL图像 - img = Image.frombytes("RGB", sct_img.size, sct_img.bgra, "raw", "BGRX") - else: - # 使用PIL进行截图(备选方案) - img = ImageGrab.grab(bbox=(x, y, x + width, y + height)) - - # 压缩图像 - quality = int(self.config.get("Capture", "quality")) - buffer = BytesIO() - img.save(buffer, format="JPEG", quality=quality) - - # 转换为Base64 - img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") - - return { - "image": img_base64, - "window_rect": [x, y, width, height], - "timestamp": time.time() - } - - except Exception as e: - logger.error(f"截取游戏画面失败: {e}") - logger.error(traceback.format_exc()) - return None - - async def send_heartbeat(self): - """发送心跳包""" - if not self.ws or not self.running: - return False - - try: - await self.ws.send(json.dumps({ - "type": "heartbeat", - "timestamp": time.time(), - "client_id": self.client_id, - "game_state": { - "in_battle": self.game_state["in_battle"], - "current_map": self.game_state["current_map"] - } - })) - - # 等待心跳响应 - response_raw = await asyncio.wait_for( - self.ws.recv(), - timeout=5.0 - ) - - # 检查响应 - response = json.loads(response_raw) - if response.get("type") != "heartbeat_response": - logger.warning(f"收到非心跳响应: {response.get('type')}") - return False - - self.last_heartbeat_time = time.time() - return True - - except asyncio.TimeoutError: - logger.warning("心跳超时") - return False - except websockets.exceptions.ConnectionClosed: - logger.warning("发送心跳时连接已关闭") - return False - except Exception as e: - logger.error(f"发送心跳失败: {e}") - return False - - async def heartbeat_loop(self): - """心跳循环""" - heartbeat_interval = float(self.config.get("Connection", "heartbeat_interval", fallback="5")) - - while self.running: - if self.ws and not self.ws.closed: - success = await self.send_heartbeat() - if not success: - # 心跳失败，检查连接状态 - if time.time() - self.last_heartbeat_time > heartbeat_interval * 3: - logger.warning(f"心跳超时超过 {heartbeat_interval * 3} 秒，尝试重新连接") - if self.ws: - await self.ws.close() - self.ws = None - - await asyncio.sleep(heartbeat_interval) - - async def reconnect_loop(self): - """重连检查循环""" - while self.running: - if not self.ws or self.ws.closed: - logger.warning("WebSocket连接已断开，尝试重连...") - if await self.connect_with_retry(): - logger.info("重连成功") - else: - logger.error("重连失败") - # 不要立即停止，继续尝试 - - await asyncio.sleep(5) # 每5秒检查一次连接状态 - - async def execute_action(self, action): - """执行操作""" - try: - action_type = action.get("type") - - # 等待指定的延迟时间 - if "delay" in action: - await asyncio.sleep(action["delay"]) - - # 获取游戏窗口 - hwnd = self.get_game_window() - if hwnd is None: - return - - # 确保窗口处于前台 - if win32gui.GetForegroundWindow() != hwnd: - try: - # 尝试多种方法激活窗口 - win32gui.ShowWindow(hwnd, win32con.SW_RESTORE) # 恢复窗口（如果最小化） - win32gui.SetForegroundWindow(hwnd) # 尝试置为前台 - await asyncio.sleep(0.1) - - # 如果窗口仍然不在前台，使用更强的方法 - if win32gui.GetForegroundWindow() != hwnd: - # 获取当前激活窗口的线程和进程ID - curr_hwnd = win32gui.GetForegroundWindow() - curr_thread_id = win32process.GetWindowThreadProcessId(curr_hwnd)[0] - # 获取目标窗口的线程和进程ID - target_thread_id = win32process.GetWindowThreadProcessId(hwnd)[0] - # 附加线程输入 - win32process.AttachThreadInput(target_thread_id, curr_thread_id, True) - win32gui.SetForegroundWindow(hwnd) - win32gui.BringWindowToTop(hwnd) - win32process.AttachThreadInput(target_thread_id, curr_thread_id, False) - await asyncio.sleep(0.1) - except Exception as e: - logger.error(f"激活窗口失败: {e}") - - # 获取窗口位置 - rect = win32gui.GetWindowRect(hwnd) - window_x, window_y = rect[0], rect[1] - - # 执行不同类型的操作 - if action_type == "move_to": - # 移动到指定位置 - position = action.get("position", [0, 0]) - x, y = position[0] + window_x, position[1] + window_y - - # 生成人类化的移动路径 - current_pos = win32gui.GetCursorPos() - path = self.generate_movement_path(current_pos, [x, y]) - - # 执行移动 - for point in path: - win32api.SetCursorPos((int(point[0]), int(point[1]))) - await asyncio.sleep(0.01) # 10ms延迟 - - elif action_type == "click": - # 点击指定位置 - position = action.get("position", [0, 0]) - x, y = position[0] + window_x, position[1] + window_y - - # 移动到位置 - current_pos = win32gui.GetCursorPos() - path = self.generate_movement_path(current_pos, [x, y]) - - for point in path: - win32api.SetCursorPos((int(point[0]), int(point[1]))) - await asyncio.sleep(0.01) - - # 执行点击 - mouse.click() - - elif action_type == "use_skill": - # 使用技能 - key = action.get("key", "1") - keyboard.press(key) - await asyncio.sleep(0.05) - keyboard.release(key) - - # 如果有目标位置，移动鼠标并点击 - if "target_position" in action: - pos = action["target_position"] - x, y = pos[0] + window_x, pos[1] + window_y - win32api.SetCursorPos((int(x), int(y))) - await asyncio.sleep(0.05) - mouse.click() - - elif action_type == "interact": - # 交互 - key = action.get("key", "f") - keyboard.press(key) - await asyncio.sleep(0.1) - keyboard.release(key) - - elif action_type == "move_random": - # 随机移动 - direction = action.get("direction", "right") - duration = action.get("duration", 1.0) - - # 方向键映射 - dir_keys = { - "up": "w", - "down": "s", - "left": "a", - "right": "d" - } - - key = dir_keys.get(direction, "d") - keyboard.press(key) - await asyncio.sleep(duration) - keyboard.release(key) - - elif action_type == "use_item": - # 使用物品 - key = action.get("key", "f1") - keyboard.press(key) - await asyncio.sleep(0.1) - keyboard.release(key) - - elif action_type == "stop": - # 停止所有按键 - for key in ["w", "a", "s", "d", "1", "2", "3", "4", "5", "6"]: - if keyboard.is_pressed(key): - keyboard.release(key) - - elif action_type == "type_text": - # 输入文本 - text = action.get("text", "") - if text: - keyboard.write(text, delay=0.05) - - elif action_type == "press_key_combo": - # 按下组合键 - keys = action.get("keys", []) - if keys: - for key in keys: - keyboard.press(key) - await asyncio.sleep(0.1) - for key in reversed(keys): - keyboard.release(key) - - else: - logger.warning(f"未知操作类型: {action_type}") - - # 更新游戏状态 - self.game_state["last_operation_time"] = time.time() - - except Exception as e: - logger.error(f"执行操作失败: {e}") - logger.error(traceback.format_exc()) - - def generate_movement_path(self, start_pos, end_pos, steps=None): - """生成模拟人类的鼠标移动路径""" - # 计算距离 - distance = ((end_pos[0] - start_pos[0])**2 + (end_pos[1] - start_pos[1])**2)**0.5 - - # 如果未指定步数，根据距离计算 - if steps is None: - steps = max(int(distance / 20), 5) # 每20像素一个点，至少5个点 - - # 生成基础路径 - t = np.linspace(0, 1, steps) - path = [] - - # 动态调整平滑度 - smoothness = float(self.config.get("Game", "movement_smoothness", fallback="0.8")) - - for i in range(steps): - # 基础线性插值 - x = start_pos[0] + (end_pos[0] - start_pos[0]) * t[i] - y = start_pos[1] + (end_pos[1] - start_pos[1]) * t[i] - - # 添加随机偏移（越靠近中间偏移越大） - mid_factor = 4 * t[i] * (1 - t[i]) # 在中间最大 - max_offset = distance * 0.05 * mid_factor * (1 - smoothness) # 最大偏移为距离的5%，受平滑度影响 - - offset_x = random.normalvariate(0, max_offset / 3) - offset_y = random.normalvariate(0, max_offset / 3) - - # 添加到路径 - path.append([x + offset_x, y + offset_y]) - - # 确保起点和终点准确 - path[0] = start_pos - path[-1] = end_pos - - return path - - def analyze_image(self, image_data, detection_results): - """ - 分析图像数据，更新游戏状态 - - 参数: - image_data (dict): 图像数据 - detection_results (list): 检测结果 - """ - # 更新战斗状态 - monsters_detected = False - for det in detection_results: - if det["class_name"] in ["monster", "boss"]: - monsters_detected = True - break - - self.game_state["in_battle"] = monsters_detected - - # 分析血条和蓝条 - hp_bars = [d for d in detection_results if d["class_name"] == "hp_bar"] - mp_bars = [d for d in detection_results if d["class_name"] == "mp_bar"] - - if hp_bars: - # 估算血量百分比 - self.game_state["hp_percent"] = self.estimate_bar_percent(hp_bars[0]) - - if mp_bars: - # 估算蓝量百分比 - self.game_state["mp_percent"] = self.estimate_bar_percent(mp_bars[0]) - - # 检测技能冷却 - cooldowns = [d for d in detection_results if d["class_name"] == "cooldown"] - self.game_state["cooldowns"] = {} - - for cd in cooldowns: - if "skill_id" in cd: - self.game_state["cooldowns"][cd["skill_id"]] = cd.get("remaining_time", 1.0) - - def estimate_bar_percent(self, bar_detection): - """ - 估计血条/蓝条的百分比 - - 参数: - bar_detection (dict): 条形检测结果 - - 返回: - float: 百分比值(0-100) - """ - # 这里需要根据实际情况实现 - # 临时方案：返回检测结果中的值，如果没有则返回默认值 - return bar_detection.get("percent", 100) - - async def capture_and_process_loop(self): - """截图和处理循环""" - consecutive_errors = 0 - - while self.running: - try: - # 检查WebSocket连接 - if not self.ws or self.ws.closed: - await asyncio.sleep(0.5) # 连接断开时等待 - continue - - # 获取游戏窗口 - hwnd = self.get_game_window() - if hwnd is None: - await asyncio.sleep(1.0) # 找不到窗口时等待1秒 - continue - - # 截取游戏画面 - screen_data = self.capture_game_screen(hwnd) - if screen_data is None: - await asyncio.sleep(0.5) # 截图失败时等待0.5秒 - continue - - # 准备请求数据 - request = { - "type": "image", - "request_id": f"req_{random.getrandbits(32)}", - "timestamp": time.time(), - "data": screen_data["image"], - "game_state": self.game_state, - "window_rect": screen_data["window_rect"] - } - - # 发送请求 - await self.ws.send(json.dumps(request)) - - # 等待响应 - response_raw = await asyncio.wait_for( - self.ws.recv(), - timeout=5.0 - ) - - # 处理响应 - response = json.loads(response_raw) - - if response.get("type") == "action_response": - # 获取检测结果并更新游戏状态 - if "detections" in response: - self.analyze_image(screen_data, response["detections"]) - - # 执行动作 - actions = response.get("actions", []) - - # 按优先级排序 - actions.sort(key=lambda x: x.get("execution_priority", 1.0)) - - for action in actions: - await self.execute_action(action) - - # 重置错误计数 - consecutive_errors = 0 - - elif response.get("type") == "error": - logger.error(f"服务器返回错误: {response.get('message')}") - consecutive_errors += 1 - - # 等待指定的间隔时间 - await asyncio.sleep(self.capture_interval) - - except asyncio.TimeoutError: - logger.warning("等待服务器响应超时") - consecutive_errors += 1 - except websockets.exceptions.ConnectionClosed: - logger.error("WebSocket连接已关闭") - break - except Exception as e: - logger.error(f"处理循环出错: {e}") - logger.error(traceback.format_exc()) - consecutive_errors += 1 - await asyncio.sleep(1.0) # 出错时等待1秒 - - # 如果连续错误过多，尝试重新连接 - if consecutive_errors >= 5: - logger.warning(f"连续出错 {consecutive_errors} 次，尝试重新连接") - if self.ws: - await self.ws.close() - self.ws = None - consecutive_errors = 0 - - async def run(self): - """运行客户端""" - self.running = True - - # 连接到服务器 - if not await self.connect_with_retry(): - self.running = False - return - - try: - # 创建任务 - capture_task = asyncio.create_task(self.capture_and_process_loop()) - heartbeat_task = asyncio.create_task(self.heartbeat_loop()) - reconnect_task = asyncio.create_task(self.reconnect_loop()) - - # 等待任务完成 - await asyncio.gather(capture_task, heartbeat_task, reconnect_task) - - except asyncio.CancelledError: - logger.info("客户端任务已取消") - except Exception as e: - logger.error(f"客户端运行出错: {e}") - logger.error(traceback.format_exc()) - finally: - self.running = False - if self.ws: - await self.ws.close() - logger.info("客户端已停止") - - def start(self): - """启动客户端""" - try: - # 启动心跳监控线程（备用方案，以防异步心跳失效） - self._monitor_thread = threading.Thread(target=self._monitor_connection) - self._monitor_thread.daemon = True - self._monitor_thread.start() - - # 运行主循环 - asyncio.run(self.run()) - except KeyboardInterrupt: - logger.info("用户中断，正在退出...") - except Exception as e: - logger.error(f"客户端出错: {e}") - logger.error(traceback.format_exc()) - - def _monitor_connection(self): - """监控连接的后台线程""" - while True: - try: - time.sleep(30) # 每30秒检查一次 - - if not self.running: - break - - # 检查心跳时间 - if self.last_heartbeat_time > 0 and time.time() - self.last_heartbeat_time > 60: - logger.warning("心跳超时，可能需要重连") - # 不直接重连，留给重连循环处理 - except Exception as e: - logger.error(f"连接监控线程出错: {e}") - - def stop(self): - """停止客户端""" - self.running = False - logger.info("正在停止客户端...") - -# 创建默认配置文件（如果不存在） -def ensure_config(): - if not os.path.exists(CONFIG_FILE): - config = configparser.ConfigParser() - - config["Server"] = { - "url": "wss://your-server-url:8080/ws", - "verify_ssl": "false" - } - - config["Capture"] = { - "interval": "0.5", - "quality": "70" - } - - config["Game"] = { - "window_title": "地下城与勇士", - "key_mapping": "default", - "movement_smoothness": "0.8" - } - - config["Connection"] = { - "max_retries": "5", - "retry_delay": "5", - "heartbeat_interval": "5" - } - - # 保存配置 - with open(CONFIG_FILE, "w", encoding="utf-8") as f: - config.write(f) - - print(f"已创建默认配置文件: {CONFIG_FILE}") - print("请编辑配置文件设置正确的服务器地址等信息") - -# 启动客户端 -if __name__ == "__main__": - # 确保配置文件存在 - ensure_config() - - try: - client = DNFAutoClient() - client.start() - except KeyboardInterrupt: - logger.info("用户中断，正在退出...") - except Exception as e: - logger.error(f"客户端出错: {e}") - logger.error(traceback.format_exc()) \ No newline at end of file diff --git a/dnf-client-package/config.ini b/dnf-client-package/config.ini deleted file mode 100644 index cd2b483..0000000 --- a/dnf-client-package/config.ini +++ /dev/null @@ -1,17 +0,0 @@ -[Server] -url = wss://your-server-ip:8080/ws -verify_ssl = false - -[Capture] -interval = 0.5 -quality = 70 - -[Game] -window_title = 地下城与勇士 -key_mapping = default -movement_smoothness = 0.8 - -[Connection] -max_retries = 5 -retry_delay = 5 -heartbeat_interval = 5 diff --git a/dnf-client-package/start.bat b/dnf-client-package/start.bat deleted file mode 100644 index 6d0a6cf..0000000 --- a/dnf-client-package/start.bat +++ /dev/null @@ -1,4 +0,0 @@ -@echo off -echo 正在启动DNF自动化客户端... -python client.py -pause diff --git a/dnf-client-windows.zip b/dnf-client-windows.zip deleted file mode 100644 index 3f6d7d8..0000000 Binary files a/dnf-client-windows.zip and /dev/null differ diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/weights/labels.jpg b/models/weights/labels.jpg index 9eebb7b..597d804 100644 Binary files a/models/weights/labels.jpg and b/models/weights/labels.jpg differ diff --git a/models/weights/labels_correlogram.jpg b/models/weights/labels_correlogram.jpg index 530fd43..a601d5e 100644 Binary files a/models/weights/labels_correlogram.jpg and b/models/weights/labels_correlogram.jpg differ diff --git a/models/weights/results.csv b/models/weights/results.csv index b9db8d8..78f231c 100644 --- a/models/weights/results.csv +++ b/models/weights/results.csv @@ -99,3 +99,100 @@ 97, 0.0092242, 0.0074598, 0.0032303, 0.98604, 0.98365, 0.99366, 0.9308, 0.0082706, 0.005929, 0.0022603, 0.000496, 0.000496, 0.000496 98, 0.0088322, 0.007461, 0.0025581, 0.9797, 0.9832, 0.99359, 0.92623, 0.0086755, 0.006093, 0.0023018, 0.000397, 0.000397, 0.000397 99, 0.009008, 0.0076601, 0.0027773, 0.97461, 0.98322, 0.99354, 0.93797, 0.0082771, 0.0060029, 0.0023297, 0.000298, 0.000298, 0.000298 + 0, 0.10537, 0.076466, 0.071107, 0.0057694, 0.53319, 0.039631, 0.0090939, 0.083053, 0.030588, 0.065088, 0.070833, 0.0032407, 0.0032407 + 1, 0.084226, 0.072698, 0.061562, 0.39192, 0.37511, 0.099785, 0.024765, 0.078095, 0.027365, 0.054508, 0.040768, 0.006509, 0.006509 + 2, 0.081015, 0.064917, 0.053284, 0.45499, 0.44536, 0.25051, 0.085181, 0.076592, 0.021748, 0.044283, 0.010637, 0.0097112, 0.0097112 + 3, 0.074563, 0.061782, 0.048053, 0.73241, 0.43019, 0.47748, 0.18303, 0.062812, 0.018912, 0.037823, 0.009703, 0.009703, 0.009703 + 4, 0.068283, 0.057336, 0.042833, 0.67085, 0.54173, 0.50967, 0.22036, 0.054194, 0.01771, 0.0334, 0.009703, 0.009703, 0.009703 + 5, 0.061134, 0.053866, 0.037512, 0.65146, 0.45436, 0.555, 0.26817, 0.052837, 0.016332, 0.028931, 0.009604, 0.009604, 0.009604 + 6, 0.05909, 0.053523, 0.033439, 0.57689, 0.67196, 0.67603, 0.32589, 0.047249, 0.014892, 0.025372, 0.009505, 0.009505, 0.009505 + 7, 0.056171, 0.05248, 0.031248, 0.79275, 0.60855, 0.73263, 0.391, 0.045996, 0.015203, 0.021998, 0.009406, 0.009406, 0.009406 + 8, 0.054112, 0.048681, 0.027633, 0.88722, 0.72512, 0.76046, 0.38817, 0.042352, 0.014402, 0.01973, 0.009307, 0.009307, 0.009307 + 9, 0.052162, 0.048991, 0.025344, 0.89922, 0.72099, 0.82979, 0.40421, 0.042111, 0.013725, 0.017555, 0.009208, 0.009208, 0.009208 + 10, 0.0513, 0.048124, 0.024098, 0.9543, 0.73898, 0.87995, 0.53812, 0.035614, 0.013525, 0.015445, 0.009109, 0.009109, 0.009109 + 11, 0.049729, 0.046288, 0.022249, 0.80022, 0.75284, 0.81464, 0.46789, 0.039512, 0.01382, 0.014261, 0.00901, 0.00901, 0.00901 + 12, 0.047242, 0.046023, 0.020322, 0.87033, 0.82, 0.87714, 0.53135, 0.036358, 0.013075, 0.013154, 0.008911, 0.008911, 0.008911 + 13, 0.045947, 0.049008, 0.019099, 0.79432, 0.81413, 0.85648, 0.54684, 0.036657, 0.013149, 0.011747, 0.008812, 0.008812, 0.008812 + 14, 0.044563, 0.045478, 0.017354, 0.86486, 0.82851, 0.86331, 0.58533, 0.030859, 0.012629, 0.010733, 0.008713, 0.008713, 0.008713 + 15, 0.044386, 0.045777, 0.016909, 0.87888, 0.84115, 0.88297, 0.5868, 0.030465, 0.012362, 0.010119, 0.008614, 0.008614, 0.008614 + 16, 0.044421, 0.044363, 0.015621, 0.93144, 0.82596, 0.88851, 0.60501, 0.031073, 0.012311, 0.0097153, 0.008515, 0.008515, 0.008515 + 17, 0.042487, 0.046041, 0.015011, 0.92876, 0.83356, 0.89565, 0.55739, 0.031302, 0.012821, 0.0089586, 0.008416, 0.008416, 0.008416 + 18, 0.042481, 0.041832, 0.014583, 0.84688, 0.87297, 0.88565, 0.57654, 0.032206, 0.012529, 0.008568, 0.008317, 0.008317, 0.008317 + 19, 0.041834, 0.041992, 0.014201, 0.8425, 0.87675, 0.87723, 0.56809, 0.035273, 0.012553, 0.0082376, 0.008218, 0.008218, 0.008218 + 20, 0.041529, 0.040971, 0.013827, 0.88358, 0.89188, 0.90811, 0.61107, 0.029281, 0.012575, 0.0083256, 0.008119, 0.008119, 0.008119 + 21, 0.039979, 0.041675, 0.013442, 0.90751, 0.87639, 0.88663, 0.65531, 0.028874, 0.011992, 0.0077705, 0.00802, 0.00802, 0.00802 + 22, 0.040631, 0.042681, 0.013149, 0.92106, 0.89742, 0.90632, 0.63937, 0.027697, 0.012193, 0.0073721, 0.007921, 0.007921, 0.007921 + 23, 0.039521, 0.041223, 0.012621, 0.89472, 0.8844, 0.909, 0.59444, 0.029174, 0.012005, 0.0069702, 0.007822, 0.007822, 0.007822 + 24, 0.039209, 0.040794, 0.012326, 0.94475, 0.8852, 0.91868, 0.64485, 0.029654, 0.012124, 0.0069315, 0.007723, 0.007723, 0.007723 + 25, 0.038706, 0.042495, 0.012114, 0.94535, 0.88812, 0.91316, 0.66164, 0.025027, 0.011684, 0.0064895, 0.007624, 0.007624, 0.007624 + 26, 0.038853, 0.039949, 0.011693, 0.93719, 0.88993, 0.91549, 0.64189, 0.031298, 0.011859, 0.0065789, 0.007525, 0.007525, 0.007525 + 27, 0.037325, 0.038238, 0.011581, 0.94332, 0.8866, 0.91919, 0.69192, 0.024287, 0.011155, 0.0062871, 0.007426, 0.007426, 0.007426 + 28, 0.037054, 0.039379, 0.010862, 0.94197, 0.88904, 0.91807, 0.67987, 0.027484, 0.011633, 0.0063429, 0.007327, 0.007327, 0.007327 + 29, 0.037804, 0.039336, 0.010948, 0.94615, 0.88732, 0.92419, 0.71138, 0.024992, 0.011194, 0.0060732, 0.007228, 0.007228, 0.007228 + 30, 0.035318, 0.038011, 0.010588, 0.93566, 0.90238, 0.9254, 0.69709, 0.025254, 0.01115, 0.0060346, 0.007129, 0.007129, 0.007129 + 31, 0.035848, 0.040036, 0.010657, 0.95219, 0.87909, 0.91779, 0.68068, 0.02775, 0.011805, 0.0060232, 0.00703, 0.00703, 0.00703 + 32, 0.036455, 0.03784, 0.010277, 0.9521, 0.89118, 0.91821, 0.69482, 0.02398, 0.011316, 0.006109, 0.006931, 0.006931, 0.006931 + 33, 0.035511, 0.040274, 0.01034, 0.95571, 0.89172, 0.91894, 0.67729, 0.025415, 0.011374, 0.0058415, 0.006832, 0.006832, 0.006832 + 34, 0.034818, 0.036288, 0.0098505, 0.96073, 0.90876, 0.92328, 0.73266, 0.022878, 0.010763, 0.0055736, 0.006733, 0.006733, 0.006733 + 35, 0.034432, 0.035868, 0.0099866, 0.94404, 0.89441, 0.90445, 0.6819, 0.02454, 0.010886, 0.0055561, 0.006634, 0.006634, 0.006634 + 36, 0.034021, 0.035818, 0.009415, 0.95426, 0.90633, 0.90901, 0.68886, 0.023201, 0.010798, 0.0054224, 0.006535, 0.006535, 0.006535 + 37, 0.033418, 0.03598, 0.0094498, 0.94152, 0.89777, 0.89953, 0.69698, 0.023835, 0.010837, 0.0053287, 0.006436, 0.006436, 0.006436 + 38, 0.033272, 0.036502, 0.0095944, 0.94272, 0.90222, 0.92314, 0.70942, 0.023785, 0.011131, 0.0052042, 0.006337, 0.006337, 0.006337 + 39, 0.033521, 0.035348, 0.0089882, 0.94528, 0.90316, 0.9228, 0.7345, 0.022518, 0.010903, 0.0052435, 0.006238, 0.006238, 0.006238 + 40, 0.032908, 0.035765, 0.0095038, 0.94537, 0.90443, 0.92117, 0.73349, 0.023005, 0.010825, 0.0052056, 0.006139, 0.006139, 0.006139 + 41, 0.031944, 0.03523, 0.0090051, 0.94376, 0.89426, 0.92003, 0.72121, 0.021874, 0.010633, 0.0051577, 0.00604, 0.00604, 0.00604 + 42, 0.032284, 0.034942, 0.0085232, 0.94648, 0.89959, 0.92236, 0.72606, 0.021799, 0.010527, 0.0051405, 0.005941, 0.005941, 0.005941 + 43, 0.03217, 0.035902, 0.0087585, 0.9491, 0.9019, 0.92087, 0.73303, 0.021869, 0.01047, 0.0051244, 0.005842, 0.005842, 0.005842 + 44, 0.031443, 0.033035, 0.0085704, 0.94745, 0.90156, 0.92477, 0.7546, 0.020727, 0.01034, 0.0049707, 0.005743, 0.005743, 0.005743 + 45, 0.031114, 0.036149, 0.0084183, 0.94662, 0.90115, 0.92202, 0.72985, 0.021511, 0.010381, 0.0051055, 0.005644, 0.005644, 0.005644 + 46, 0.030806, 0.033467, 0.0085668, 0.94773, 0.9066, 0.91584, 0.7378, 0.021743, 0.010383, 0.0049932, 0.005545, 0.005545, 0.005545 + 47, 0.03045, 0.035342, 0.0088827, 0.94675, 0.89267, 0.91217, 0.72865, 0.021152, 0.010576, 0.0050758, 0.005446, 0.005446, 0.005446 + 48, 0.030546, 0.034058, 0.0083615, 0.9474, 0.89399, 0.90736, 0.72605, 0.022204, 0.010774, 0.0049421, 0.005347, 0.005347, 0.005347 + 49, 0.030783, 0.035215, 0.0082594, 0.94012, 0.88191, 0.91291, 0.76995, 0.019367, 0.010101, 0.0049213, 0.005248, 0.005248, 0.005248 + 50, 0.030272, 0.031735, 0.0081784, 0.94596, 0.89436, 0.91247, 0.72832, 0.022859, 0.010223, 0.0047657, 0.005149, 0.005149, 0.005149 + 51, 0.030048, 0.033497, 0.0078879, 0.9442, 0.90128, 0.91924, 0.77297, 0.018489, 0.0096083, 0.0046946, 0.00505, 0.00505, 0.00505 + 52, 0.029478, 0.034815, 0.0079617, 0.94805, 0.90482, 0.91374, 0.73906, 0.021018, 0.010099, 0.0047801, 0.004951, 0.004951, 0.004951 + 53, 0.029622, 0.033132, 0.0073636, 0.94613, 0.90688, 0.90788, 0.76831, 0.018532, 0.0098147, 0.0046803, 0.004852, 0.004852, 0.004852 + 54, 0.029849, 0.034944, 0.0074608, 0.95061, 0.90497, 0.90431, 0.73547, 0.020199, 0.010039, 0.0046062, 0.004753, 0.004753, 0.004753 + 55, 0.028862, 0.030966, 0.0074558, 0.95399, 0.90448, 0.92083, 0.77247, 0.017624, 0.0094819, 0.0045197, 0.004654, 0.004654, 0.004654 + 56, 0.028374, 0.031805, 0.0075731, 0.94923, 0.9021, 0.91864, 0.7448, 0.021008, 0.0099044, 0.0046206, 0.004555, 0.004555, 0.004555 + 57, 0.028933, 0.033019, 0.0077383, 0.94261, 0.90025, 0.91444, 0.76969, 0.018055, 0.0092425, 0.0044035, 0.004456, 0.004456, 0.004456 + 58, 0.02764, 0.031741, 0.0071853, 0.9316, 0.90703, 0.90207, 0.7154, 0.019898, 0.0099125, 0.0043523, 0.004357, 0.004357, 0.004357 + 59, 0.027372, 0.032124, 0.0071181, 0.93235, 0.90338, 0.91095, 0.78268, 0.017186, 0.0096073, 0.0043241, 0.004258, 0.004258, 0.004258 + 60, 0.027831, 0.03038, 0.0073338, 0.93878, 0.90209, 0.91096, 0.74039, 0.02007, 0.0099401, 0.0042286, 0.004159, 0.004159, 0.004159 + 61, 0.028398, 0.031714, 0.0072932, 0.9362, 0.90344, 0.92164, 0.79353, 0.017023, 0.0093412, 0.0041499, 0.00406, 0.00406, 0.00406 + 62, 0.02799, 0.034212, 0.0075848, 0.92984, 0.90388, 0.90626, 0.74251, 0.01955, 0.0094862, 0.0041463, 0.003961, 0.003961, 0.003961 + 63, 0.026864, 0.029966, 0.0072165, 0.9445, 0.90611, 0.9103, 0.75938, 0.018653, 0.0093565, 0.0041413, 0.003862, 0.003862, 0.003862 + 64, 0.026335, 0.029561, 0.0066826, 0.93036, 0.90621, 0.92241, 0.78571, 0.017073, 0.0095363, 0.0040866, 0.003763, 0.003763, 0.003763 + 65, 0.026865, 0.032126, 0.0069433, 0.93762, 0.8962, 0.91818, 0.76424, 0.019093, 0.0098457, 0.0040423, 0.003664, 0.003664, 0.003664 + 66, 0.025874, 0.029991, 0.0067089, 0.93369, 0.90285, 0.91089, 0.75601, 0.01744, 0.0095415, 0.0040609, 0.003565, 0.003565, 0.003565 + 67, 0.025686, 0.029558, 0.0068254, 0.94118, 0.89702, 0.93001, 0.80512, 0.016593, 0.0091715, 0.0040623, 0.003466, 0.003466, 0.003466 + 68, 0.026146, 0.0306, 0.006672, 0.9435, 0.89793, 0.92328, 0.78872, 0.016968, 0.0094819, 0.0039386, 0.003367, 0.003367, 0.003367 + 69, 0.02575, 0.03149, 0.0066763, 0.95011, 0.89324, 0.92115, 0.77461, 0.017017, 0.00915, 0.0038871, 0.003268, 0.003268, 0.003268 + 70, 0.025257, 0.029951, 0.0062569, 0.94152, 0.90354, 0.93186, 0.79316, 0.016541, 0.0089645, 0.0037884, 0.003169, 0.003169, 0.003169 + 71, 0.025736, 0.031298, 0.0063836, 0.94456, 0.89936, 0.98277, 0.85375, 0.015854, 0.0091391, 0.0037851, 0.00307, 0.00307, 0.00307 + 72, 0.025147, 0.029945, 0.0065308, 0.94031, 0.90261, 0.98337, 0.80994, 0.016953, 0.0095129, 0.0037837, 0.002971, 0.002971, 0.002971 + 73, 0.025429, 0.029684, 0.0067952, 0.9484, 0.89445, 0.97873, 0.85639, 0.015808, 0.0090361, 0.0037594, 0.002872, 0.002872, 0.002872 + 74, 0.0244, 0.02808, 0.006261, 0.9484, 0.89159, 0.96903, 0.8146, 0.01757, 0.0091019, 0.0037436, 0.002773, 0.002773, 0.002773 + 75, 0.024458, 0.029165, 0.0064301, 0.9414, 0.90268, 0.96465, 0.82351, 0.016609, 0.0093188, 0.0037172, 0.002674, 0.002674, 0.002674 + 76, 0.024491, 0.028867, 0.0065167, 0.93885, 0.90952, 0.96421, 0.82762, 0.015425, 0.0091639, 0.0037354, 0.002575, 0.002575, 0.002575 + 77, 0.024234, 0.030084, 0.0061172, 0.92641, 0.92039, 0.97935, 0.84484, 0.015876, 0.0093784, 0.0036685, 0.002476, 0.002476, 0.002476 + 78, 0.023912, 0.027955, 0.0057952, 0.92526, 0.92039, 0.97048, 0.83617, 0.016671, 0.0094657, 0.0036557, 0.002377, 0.002377, 0.002377 + 79, 0.023572, 0.028896, 0.0056073, 0.94601, 0.89761, 0.96783, 0.83554, 0.015535, 0.0090661, 0.0036031, 0.002278, 0.002278, 0.002278 + 80, 0.023433, 0.027542, 0.0060494, 0.95198, 0.89309, 0.96861, 0.83564, 0.01569, 0.0090857, 0.0035759, 0.002179, 0.002179, 0.002179 + 81, 0.023474, 0.029614, 0.0059355, 0.956, 0.89106, 0.96929, 0.83867, 0.015587, 0.0091519, 0.0035584, 0.00208, 0.00208, 0.00208 + 82, 0.022964, 0.028318, 0.0058851, 0.95519, 0.89445, 0.97281, 0.84384, 0.01506, 0.0089855, 0.003513, 0.001981, 0.001981, 0.001981 + 83, 0.023384, 0.028354, 0.0057465, 0.95494, 0.89426, 0.98393, 0.86829, 0.015179, 0.0090079, 0.0034722, 0.001882, 0.001882, 0.001882 + 84, 0.022817, 0.027814, 0.0059098, 0.95244, 0.89233, 0.97627, 0.84317, 0.015875, 0.0092964, 0.0034103, 0.001783, 0.001783, 0.001783 + 85, 0.023342, 0.027297, 0.0060606, 0.95319, 0.89174, 0.97141, 0.84056, 0.015359, 0.0091238, 0.0033967, 0.001684, 0.001684, 0.001684 + 86, 0.022191, 0.027358, 0.005542, 0.9539, 0.88966, 0.97417, 0.85334, 0.015129, 0.0091815, 0.0034193, 0.001585, 0.001585, 0.001585 + 87, 0.022573, 0.027708, 0.0056441, 0.95191, 0.89295, 0.97479, 0.84976, 0.015716, 0.0091705, 0.0033853, 0.001486, 0.001486, 0.001486 + 88, 0.022409, 0.027443, 0.0058222, 0.92751, 0.9028, 0.96398, 0.84426, 0.015372, 0.008944, 0.0033563, 0.001387, 0.001387, 0.001387 + 89, 0.021741, 0.026495, 0.0053773, 0.94905, 0.89763, 0.97, 0.84371, 0.014766, 0.0090199, 0.0033002, 0.001288, 0.001288, 0.001288 + 90, 0.021996, 0.028446, 0.0058893, 0.93917, 0.88982, 0.96402, 0.8402, 0.01475, 0.0089879, 0.0032734, 0.001189, 0.001189, 0.001189 + 91, 0.022016, 0.027665, 0.0055442, 0.94716, 0.90265, 0.97066, 0.849, 0.01469, 0.0089841, 0.0033016, 0.00109, 0.00109, 0.00109 + 92, 0.021633, 0.027603, 0.0055049, 0.9475, 0.90336, 0.96807, 0.84835, 0.014596, 0.0091195, 0.0032879, 0.000991, 0.000991, 0.000991 + 93, 0.021683, 0.026557, 0.0052333, 0.94842, 0.90211, 0.97055, 0.85609, 0.014893, 0.0090179, 0.0032596, 0.000892, 0.000892, 0.000892 + 94, 0.021309, 0.025954, 0.0053442, 0.94614, 0.90127, 0.97183, 0.85953, 0.014493, 0.0088654, 0.0032242, 0.000793, 0.000793, 0.000793 + 95, 0.021511, 0.026043, 0.0054107, 0.94789, 0.90206, 0.96895, 0.86042, 0.014186, 0.008769, 0.003192, 0.000694, 0.000694, 0.000694 + 96, 0.020815, 0.024904, 0.0052537, 0.94738, 0.90358, 0.9677, 0.84651, 0.01453, 0.0088882, 0.0031793, 0.000595, 0.000595, 0.000595 diff --git a/models/weights/train_batch0.jpg b/models/weights/train_batch0.jpg index 3e1fa75..370780b 100644 Binary files a/models/weights/train_batch0.jpg and b/models/weights/train_batch0.jpg differ diff --git a/models/weights/train_batch1.jpg b/models/weights/train_batch1.jpg index 43a6391..5ca084d 100644 Binary files a/models/weights/train_batch1.jpg and b/models/weights/train_batch1.jpg differ diff --git a/models/weights/train_batch2.jpg b/models/weights/train_batch2.jpg index 58af309..b2c3326 100644 Binary files a/models/weights/train_batch2.jpg and b/models/weights/train_batch2.jpg differ diff --git a/models/yolo_model.py b/models/yolo_model.py index d94cdb0..5095d3e 100644 --- a/models/yolo_model.py +++ b/models/yolo_model.py @@ -2,580 +2,467 @@ # -*- coding: utf-8 -*- """ -YOLO模型封装，提供图像识别功能 -优化版 - 增强性能和模型缓存管理 -修复版 - 解决模型加载依赖问题 +YOLO模型实现，负责目标检测和识别 +支持多种YOLO版本：YOLOv5、YOLOv8 +优化版 - 支持GPU加速和模型优化 """ import os -import sys # 添加缺失的sys模块导入 -import logging -import torch -import yaml +import sys import time +import logging import numpy as np -import subprocess +from pathlib import Path +import torch +import cv2 from PIL import Image +import statistics from collections import deque -from config.settings import MODEL, BASE_DIR +from config.settings import MODEL logger = logging.getLogger("DNFAutoCloud") class YOLOModel: - """YOLO模型封装类 - 优化版""" + """YOLO模型类 - 优化版""" def __init__(self): """初始化YOLO模型""" - self.device = MODEL.get("device", "cuda" if torch.cuda.is_available() else "cpu") + self.model = None + self.device = MODEL.get("device", "cpu") self.conf_threshold = MODEL.get("conf_threshold", 0.5) self.iou_threshold = MODEL.get("iou_threshold", 0.45) + self.img_size = MODEL.get("img_size", 640) + self.half_precision = MODEL.get("half_precision", False) + self.class_names = self._load_class_names() - # 性能监控 + # 性能追踪 self.inference_times = deque(maxlen=100) - self.batch_size = MODEL.get("batch_size", 1) - self.use_half_precision = MODEL.get("half_precision", True) and self.device == "cuda" - logger.info(f"正在加载YOLO模型: {MODEL.get('name', 'unknown')}") - logger.info(f"模型权重路径: {MODEL.get('weights', 'unknown')}") - logger.info(f"设备: {self.device}, 半精度: {self.use_half_precision}") - - # 检查权重文件是否存在 - weights_path = MODEL.get("weights", "") - if not os.path.exists(weights_path): - # 尝试从备选路径加载 - alt_paths = [ - os.path.join(BASE_DIR, "models", "weights", "best.pt"), - os.path.join(BASE_DIR, "models", "best.pt"), - os.path.join(BASE_DIR, "yolov5", "runs", "train", "exp", "weights", "best.pt") - ] + # 初始化模型 + self._initialize_model() + + def _load_class_names(self): + """加载类别名称""" + try: + class_names_path = MODEL.get("class_names", "") + if not class_names_path or not os.path.exists(class_names_path): + logger.warning(f"类别名称文件不存在: {class_names_path}，使用默认类别") + return [ + "monster", "boss", "door", "item", "npc", "player", + "hp_bar", "mp_bar", "skill_ready", "cooldown" + ] - for path in alt_paths: - if os.path.exists(path): - weights_path = path - logger.info(f"使用备选模型权重: {weights_path}") - break + # 尝试从YAML加载 + try: + import yaml + with open(class_names_path, 'r', encoding='utf-8') as f: + data = yaml.safe_load(f) + + if isinstance(data, dict) and "names" in data: + return data["names"] + elif "nc" in data and isinstance(data.get("names", []), list): + return data["names"] + except: + # 如果YAML加载失败，尝试直接读取为文本 + with open(class_names_path, 'r', encoding='utf-8') as f: + class_names = [line.strip() for line in f.readlines() if line.strip()] + return class_names + + except Exception as e: + logger.error(f"加载类别名称时出错: {e}") + return [f"class_{i}" for i in range(10)] # 默认类别名 + + def _initialize_model(self): + """初始化YOLO模型""" + try: + # 获取模型路径 + weights_path = MODEL.get("weights", "") + if not weights_path: + raise ValueError("未指定模型权重路径") if not os.path.exists(weights_path): - raise FileNotFoundError(f"找不到模型权重文件: {weights_path}") - - # 加载模型 - try: - if MODEL.get("engine", "") == "onnx": - self._load_onnx_model() - else: - # 尝试不同的加载方法 - try: - self._load_yolov5_custom() - except Exception as e: - logger.warning(f"使用torch.hub加载YOLOv5模型失败: {e}") - logger.info("尝试备选方法加载模型...") - self._load_yolov5_manually() + raise FileNotFoundError(f"模型权重文件不存在: {weights_path}") - logger.info("YOLO模型加载成功") + logger.info(f"正在加载YOLO模型: {weights_path}") + + # 确定模型引擎类型 + engine = MODEL.get("engine", "pytorch").lower() + + if engine == "onnx": + # 使用ONNX模型 + if not weights_path.endswith(".onnx"): + weights_path = weights_path + ".onnx" + + self._initialize_onnx_model(weights_path) + else: + # 默认使用PyTorch模型 + self._initialize_pytorch_model(weights_path) + + logger.info(f"YOLO模型加载成功，运行于 {self.device} 设备") except Exception as e: - logger.error(f"加载YOLO模型失败: {e}") + logger.error(f"初始化YOLO模型失败: {e}") raise - def _load_yolov5_custom(self): - """使用torch.hub加载YOLOv5模型""" + def _initialize_pytorch_model(self, weights_path): + """初始化PyTorch YOLO模型""" try: - # 尝试从本地加载YOLOv5 - yolov5_dir = os.path.join(BASE_DIR, "tools", "yolov5") - if os.path.exists(yolov5_dir): - logger.info(f"从本地目录加载YOLOv5: {yolov5_dir}") - sys.path.insert(0, yolov5_dir) - + # 检查设备可用性 + if self.device.startswith("cuda") and not torch.cuda.is_available(): + logger.warning("CUDA不可用，切换到CPU模式") + self.device = "cpu" + self.half_precision = False + + # 设置设备 + device = torch.device(self.device) + + # 确定YOLO版本并加载 + model_name = MODEL.get("name", "").lower() + + if "yolov8" in model_name or weights_path.endswith("v8.pt"): + # YOLOv8 try: - # 尝试直接导入YOLOv5模块 - from models.common import DetectMultiBackend - from utils.torch_utils import select_device - from utils.general import check_img_size, non_max_suppression, scale_coords + from ultralytics import YOLO + self.model = YOLO(weights_path) + self.model_type = "yolov8" + logger.info("已加载YOLOv8模型") + except ImportError: + logger.error("无法导入ultralytics，请安装: pip install ultralytics") + raise + else: + # YOLOv5（默认） + try: + sys.path.append(os.path.join(os.path.dirname(__file__), "../tools/yolov5")) + import torch # 加载模型 - self.model = DetectMultiBackend(MODEL.get("weights"), device=self.device) - self.stride = self.model.stride - self.names = self.model.names + self.model = torch.hub.load('ultralytics/yolov5', 'custom', + path=weights_path, device=device) # 设置参数 - self.imgsz = check_img_size((640, 640), s=self.stride) + self.model.conf = self.conf_threshold + self.model.iou = self.iou_threshold + self.model.classes = None # 检测所有类别 + self.model.max_det = 100 # 最大检测数量 - # 使用半精度 - if self.use_half_precision: + # 如果启用半精度且支持 + if self.half_precision and self.device != "cpu": self.model.half() - # 预热模型 - self._warmup_model() + self.model_type = "yolov5" + logger.info("已加载YOLOv5模型") - return - except ImportError as e: - logger.warning(f"无法直接导入YOLOv5模块: {e}") - - # 尝试从torch.hub加载 - logger.info("尝试从torch.hub加载YOLOv5模型") - self.model = torch.hub.load( - 'ultralytics/yolov5', - 'custom', - path=MODEL.get("weights"), - device=self.device, - force_reload=True - ) - - # 设置模型参数 - self.model.conf = self.conf_threshold - self.model.iou = self.iou_threshold - - # 使用半精度 - if self.use_half_precision: - self.model.half() + except ImportError: + logger.error("无法导入torch或YOLOv5，请确保已正确安装") + raise + except Exception as e: + logger.error(f"加载YOLOv5模型时出错: {e}") + raise # 预热模型 self._warmup_model() except Exception as e: - logger.error(f"加载YOLOv5模型失败: {e}") + logger.error(f"初始化PyTorch模型失败: {e}") raise - def _load_yolov5_manually(self): - """手动加载YOLOv5模型（不依赖torch.hub）""" - try: - # 检查YOLOv5目录是否存在，不存在则克隆 - yolov5_dir = os.path.join(BASE_DIR, "tools", "yolov5") - if not os.path.exists(yolov5_dir): - logger.info("YOLOv5目录不存在，正在克隆仓库...") - os.makedirs(os.path.dirname(yolov5_dir), exist_ok=True) - - # 克隆YOLOv5仓库 - subprocess.run( - ["git", "clone", "https://github.com/ultralytics/yolov5.git", yolov5_dir], - check=True - ) - - # 添加YOLOv5目录到系统路径 - sys.path.insert(0, yolov5_dir) - - try: - # 尝试导入YOLOv5模块 - from models.common import DetectMultiBackend - from utils.torch_utils import select_device - from utils.general import check_img_size, non_max_suppression, scale_coords - - # 加载模型 - self.model = DetectMultiBackend(MODEL.get("weights"), device=self.device) - self.stride = self.model.stride - self.names = self.model.names - - # 设置参数 - self.imgsz = check_img_size((640, 640), s=self.stride) - - # 使用半精度 - if self.use_half_precision: - self.model.half() - - # 保存必要的函数 - self.non_max_suppression = non_max_suppression - self.scale_coords = scale_coords - - # 预热模型 - dummy_img = torch.zeros((1, 3, self.imgsz[0], self.imgsz[1]), device=self.device) - if self.use_half_precision: - dummy_img = dummy_img.half() - - with torch.no_grad(): - for _ in range(2): - self.model(dummy_img) - - logger.info("手动加载YOLOv5模型成功") - - except ImportError as e: - logger.error(f"导入YOLOv5模块失败: {e}") - # 尝试使用更简单的模型 - self._load_fallback_model() - - except Exception as e: - logger.error(f"手动加载YOLOv5模型失败: {e}") - # 尝试使用更简单的模型 - self._load_fallback_model() - - def _load_fallback_model(self): - """加载备用模型（使用PyTorch内置模型）""" - logger.info("尝试加载备用模型 (PyTorch YOLO)") - - try: - # 使用PyTorch的预训练模型 - from torchvision.models.detection import fasterrcnn_resnet50_fpn - - self.model = fasterrcnn_resnet50_fpn(pretrained=True) - self.model.to(self.device) - self.model.eval() - - # 备用模型的类别 - self.names = [ - 'background', 'monster', 'boss', 'door', 'item', 'npc', 'player', - 'hp_bar', 'mp_bar', 'skill_ready', 'cooldown' - ] - - # 标记使用备用模型 - self.using_fallback = True - - logger.info("备用模型加载成功") - - except Exception as e: - logger.error(f"加载备用模型失败: {e}") - raise - - def _load_onnx_model(self): - """加载ONNX版YOLO模型""" + def _initialize_onnx_model(self, weights_path): + """初始化ONNX YOLO模型""" try: import onnxruntime as ort - # 设置ONNX运行时参数 - if self.device == "cuda": + # 设置ONNX运行时选项 + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + + # 选择运行设备 + if self.device.startswith("cuda") and "CUDAExecutionProvider" in ort.get_available_providers(): providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] else: providers = ['CPUExecutionProvider'] - - # 检查ONNX模型文件 - onnx_path = MODEL.get("weights", "") - if not onnx_path.endswith('.onnx'): - onnx_path = onnx_path + '.onnx' - - if not os.path.exists(onnx_path): - logger.warning(f"找不到ONNX模型: {onnx_path}") - logger.info("尝试导出PyTorch模型到ONNX格式...") - - # 尝试加载PyTorch模型并导出为ONNX - self._load_yolov5_custom() - self._export_to_onnx(onnx_path) + self.device = "cpu" # 创建ONNX会话 - self.onnx_session = ort.InferenceSession(onnx_path, providers=providers) + self.model = ort.InferenceSession(weights_path, sess_options=sess_options, + providers=providers) - # 获取模型输入输出名称 - self.input_name = self.onnx_session.get_inputs()[0].name - self.output_names = [output.name for output in self.onnx_session.get_outputs()] + # 获取模型输入输出信息 + self.input_names = [input.name for input in self.model.get_inputs()] + self.output_names = [output.name for output in self.model.get_outputs()] - # 加载类别名称 - if os.path.exists(MODEL.get("class_names", "")): - with open(MODEL.get("class_names"), "r") as f: - self.class_names = yaml.safe_load(f) - else: - self.class_names = ["object"] - - # 标记使用ONNX - self.using_onnx = True + self.model_type = "onnx" + logger.info(f"已加载ONNX模型，提供者: {self.model.get_providers()}") # 预热模型 - self._warmup_onnx_model() + dummy_input = np.random.rand(1, 3, self.img_size, self.img_size).astype(np.float32) + input_feed = {self.input_names[0]: dummy_input} + self.model.run(self.output_names, input_feed) except ImportError: - logger.error("缺少ONNX运行时依赖，请安装onnxruntime或onnxruntime-gpu") - logger.info("尝试使用PyTorch模型代替...") - self._load_yolov5_custom() + logger.error("无法导入onnxruntime，请安装: pip install onnxruntime-gpu 或 onnxruntime") + raise except Exception as e: - logger.error(f"加载ONNX模型失败: {e}") - logger.info("尝试使用PyTorch模型代替...") - self._load_yolov5_custom() - - def _export_to_onnx(self, onnx_path): - """将PyTorch模型导出为ONNX格式""" - try: - import torch.onnx - - # 准备导出 - dummy_input = torch.randn(1, 3, 640, 640, device=self.device) - if self.use_half_precision: - dummy_input = dummy_input.half() - - # 导出ONNX - torch.onnx.export( - self.model, - dummy_input, - onnx_path, - verbose=False, - opset_version=12, - input_names=['images'], - output_names=['output'] - ) - - logger.info(f"PyTorch模型成功导出为ONNX格式: {onnx_path}") - - except Exception as e: - logger.error(f"导出ONNX模型失败: {e}") + logger.error(f"初始化ONNX模型失败: {e}") raise def _warmup_model(self): - """预热模型（运行空白图像以初始化）""" - dummy_img = torch.zeros((1, 3, 640, 640), device=self.device) - if self.use_half_precision: - dummy_img = dummy_img.half() - - # 进行多次推理预热 - with torch.no_grad(): - for _ in range(2): - self.model(dummy_img) - - logger.info("模型预热完成") - - def _warmup_onnx_model(self): - """预热ONNX模型""" - dummy_img = np.zeros((1, 3, 640, 640), dtype=np.float32) - - # 进行多次推理预热 - for _ in range(2): - self.onnx_session.run(self.output_names, {self.input_name: dummy_img}) - - logger.info("ONNX模型预热完成") + """预热模型，避免首次推理延迟""" + try: + logger.info("预热模型...") + + if self.model_type == "yolov8": + # YOLOv8预热 + dummy_image = np.random.randint(0, 255, (self.img_size, self.img_size, 3), dtype=np.uint8) + self.model(dummy_image, verbose=False) + + elif self.model_type == "yolov5": + # YOLOv5预热 + dummy_image = torch.zeros((1, 3, self.img_size, self.img_size), device=self.device) + if self.half_precision and self.device != "cpu": + dummy_image = dummy_image.half() + + self.model(dummy_image) + + logger.info("模型预热完成") + + except Exception as e: + logger.warning(f"模型预热出错: {e}") def detect(self, image): """ - 对图像进行目标检测 + 执行目标检测 参数: image (PIL.Image): 输入图像 返回: - list: 检测结果，每个结果包含边界框、类别和置信度 + list: 检测结果列表 """ try: start_time = time.time() - # 根据模型类型选择检测方法 - if hasattr(self, 'using_onnx') and self.using_onnx: + # 检查模型是否已初始化 + if self.model is None: + logger.error("模型未初始化") + return [] + + # 选择相应的检测方法 + if self.model_type == "yolov8": + detections = self._detect_yolov8(image) + elif self.model_type == "yolov5": + detections = self._detect_yolov5(image) + elif self.model_type == "onnx": detections = self._detect_onnx(image) - elif hasattr(self, 'using_fallback') and self.using_fallback: - detections = self._detect_fallback(image) else: - detections = self._detect_pytorch(image) + logger.error(f"不支持的模型类型: {self.model_type}") + return [] # 记录推理时间 inference_time = time.time() - start_time self.inference_times.append(inference_time) - # 计算平均推理时间 - avg_time = sum(self.inference_times) / len(self.inference_times) - - if len(self.inference_times) % 10 == 0: - logger.debug(f"平均推理时间: {avg_time:.3f}秒, 当前: {inference_time:.3f}秒") - return detections except Exception as e: logger.error(f"目标检测出错: {e}") return [] - def _detect_pytorch(self, image): - """使用PyTorch模型进行检测""" - # 在GPU上进行推理 - with torch.no_grad(): - # 处理不同的模型接口 - if hasattr(self, 'stride') and hasattr(self, 'names'): - # 自定义加载的YOLOv5 - # 转换图像 - img = self._prepare_image_custom(image) - - # 推理 - output = self.model(img) - - # 处理输出 - pred = self.non_max_suppression(output[0], self.conf_threshold, self.iou_threshold) - - # 解析结果 - detections = [] - for det in pred[0]: - x1, y1, x2, y2, conf, cls = det - detections.append({ - 'bbox': [float(x1), float(y1), float(x2), float(y2)], - 'confidence': float(conf), - 'class_id': int(cls), - 'class_name': self.names[int(cls)] - }) + def _detect_yolov5(self, image): + """使用YOLOv5模型进行检测""" + # 转换PIL图像为所需格式 + if isinstance(image, Image.Image): + # PIL图像直接传给模型 + results = self.model(image, size=self.img_size) + else: + # 转换numpy数组为PIL图像 + image = Image.fromarray(np.array(image)) + results = self.model(image, size=self.img_size) + + # 提取结果 + predictions = results.xyxy[0].cpu().numpy() # xyxy格式，(n, 6) - x1, y1, x2, y2, conf, cls + + # 转换为标准格式 + detections = [] + for pred in predictions: + x1, y1, x2, y2, conf, cls_id = pred + cls_id = int(cls_id) + + # 确保类别ID在范围内 + if cls_id < len(self.class_names): + class_name = self.class_names[cls_id] else: - # 标准torch.hub加载的模型 - results = self.model(image) - - # 处理结果 - detections = [] - for pred in results.xyxy[0].cpu().numpy(): - x1, y1, x2, y2, conf, cls = pred - detections.append({ - 'bbox': [float(x1), float(y1), float(x2), float(y2)], - 'confidence': float(conf), - 'class_id': int(cls), - 'class_name': self.model.names[int(cls)] - }) + class_name = f"class_{cls_id}" + + detections.append({ + "bbox": [float(x1), float(y1), float(x2), float(y2)], + "confidence": float(conf), + "class_id": cls_id, + "class_name": class_name + }) return detections - def _detect_fallback(self, image): - """使用备用模型进行检测""" - # 预处理图像 - img = self._prepare_image_pytorch(image) + def _detect_yolov8(self, image): + """使用YOLOv8模型进行检测""" + # 转换PIL图像为所需格式 + if isinstance(image, Image.Image): + # 转换为numpy数组 + image_np = np.array(image) + else: + image_np = np.array(image) - # 进行推理 - with torch.no_grad(): - predictions = self.model(img) + # 执行推理 + results = self.model( + image_np, + conf=self.conf_threshold, + iou=self.iou_threshold, + verbose=False + ) - # 处理结果 + # 提取结果 detections = [] - for i, prediction in enumerate(predictions[0]['boxes']): - score = predictions[0]['scores'][i].item() - if score > self.conf_threshold: - x1, y1, x2, y2 = prediction.tolist() - class_id = predictions[0]['labels'][i].item() + + for result in results: + # 获取边界框 + boxes = result.boxes + + for i in range(len(boxes)): + # 提取坐标，conf和类别 + box = boxes[i] + x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() + conf = float(box.conf) + cls_id = int(box.cls) - # 映射torchvision的COCO类别到我们的类别 - class_name = self.names[min(class_id, len(self.names) - 1)] + # 确保类别ID在范围内 + if cls_id < len(self.class_names): + class_name = self.class_names[cls_id] + else: + class_name = f"class_{cls_id}" detections.append({ - 'bbox': [float(x1), float(y1), float(x2), float(y2)], - 'confidence': float(score), - 'class_id': class_id, - 'class_name': class_name + "bbox": [float(x1), float(y1), float(x2), float(y2)], + "confidence": conf, + "class_id": cls_id, + "class_name": class_name }) return detections def _detect_onnx(self, image): """使用ONNX模型进行检测""" - # 预处理图像 - input_tensor = self._prepare_image_onnx(image) + # 准备输入图像 + if isinstance(image, Image.Image): + # 转换PIL图像为numpy数组 + img = np.array(image.resize((self.img_size, self.img_size))) + else: + # 转换任何其他格式为numpy数组 + img = cv2.resize(np.array(image), (self.img_size, self.img_size)) - # 运行推理 - outputs = self.onnx_session.run(self.output_names, {self.input_name: input_tensor}) + # 图像预处理 + img = img.transpose(2, 0, 1) # HWC -> CHW + img = np.expand_dims(img, axis=0) # 添加batch维度 + img = img / 255.0 # 归一化到[0,1] + img = img.astype(np.float32) # 转换为float32 - # 解析输出（根据ONNX模型的输出格式可能需要调整） - # 假设输出格式为 [batch_id, x1, y1, x2, y2, conf, class_id] - predictions = outputs[0] + # 执行推理 + input_feed = {self.input_names[0]: img} + outputs = self.model.run(self.output_names, input_feed) - # 过滤低置信度预测 - mask = predictions[:, 5] > self.conf_threshold - filtered_preds = predictions[mask] + # 后处理 + # 注意：实际后处理可能因模型而异 + # 这里假设输出为[batch, num_detections, 6]格式，其中6为[x1, y1, x2, y2, conf, cls] + predictions = outputs[0] # 假设第一个输出是检测结果 - # 组织检测结果 + # 转换为标准格式 detections = [] - for pred in filtered_preds: - x1, y1, x2, y2, conf, cls = pred[1:7] - cls_id = int(cls) - - detections.append({ - 'bbox': [float(x1), float(y1), float(x2), float(y2)], - 'confidence': float(conf), - 'class_id': cls_id, - 'class_name': self.class_names[cls_id] if cls_id < len(self.class_names) else "unknown" - }) + + # 提取有效检测 + valid_preds = predictions[0] # 取第一个batch + + for pred in valid_preds: + if len(pred) >= 6: # 检查是否有6个元素 + x1, y1, x2, y2, conf, cls_id = pred[:6] + + # 过滤低置信度检测 + if conf < self.conf_threshold: + continue + + cls_id = int(cls_id) + + # 确保类别ID在范围内 + if cls_id < len(self.class_names): + class_name = self.class_names[cls_id] + else: + class_name = f"class_{cls_id}" + + # 将坐标从[0,1]缩放到原始图像尺寸 + if max(x1, y1, x2, y2) <= 1.0: + # 坐标是归一化的，需要缩放到原始尺寸 + if hasattr(image, 'size'): + orig_w, orig_h = image.size + else: + orig_h, orig_w = image.shape[:2] + + x1 *= orig_w + x2 *= orig_w + y1 *= orig_h + y2 *= orig_h + + detections.append({ + "bbox": [float(x1), float(y1), float(x2), float(y2)], + "confidence": float(conf), + "class_id": cls_id, + "class_name": class_name + }) return detections - def _prepare_image_custom(self, image): - """为自定义加载的YOLOv5模型准备图像""" - # 转换PIL图像为numpy数组 - img = np.array(image) - - # 调整大小 - img = self._letterbox(img, new_shape=self.imgsz)[0] - - # BGR转RGB - img = img[:, :, ::-1].transpose(2, 0, 1) - img = np.ascontiguousarray(img) - - # 转换为PyTorch张量 - img = torch.from_numpy(img).to(self.device) - img = img.half() if self.use_half_precision else img.float() - img /= 255.0 - - # 增加批次维度 - if img.ndimension() == 3: - img = img.unsqueeze(0) - - return img - - def _prepare_image_pytorch(self, image): - """为PyTorch模型准备图像""" - # 转换PIL图像为PyTorch张量 - from torchvision import transforms - - transform = transforms.Compose([ - transforms.ToTensor(), - ]) - - img = transform(image).unsqueeze(0).to(self.device) - - return img - - def _prepare_image_onnx(self, image): - """为ONNX模型准备图像""" - # 调整图像大小 - img_size = MODEL.get("img_size", 640) - image = image.resize((img_size, img_size), Image.LANCZOS) - - # 转换为numpy数组 - img = np.array(image).astype(np.float32) / 255.0 - - # 从HWC转换为CHW格式 - img = img.transpose(2, 0, 1) - - # 添加批次维度 - img = np.expand_dims(img, axis=0) - - return img - - def _letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)): - """调整图像大小并填充（YOLOv5风格）""" - shape = img.shape[:2] # 当前形状 [高, 宽] - - # 缩放比例 (新 / 旧) - r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) - - # 计算填充 - new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) - dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] - - # 平均分配填充 - dw /= 2 - dh /= 2 - - # 调整大小 - if shape[::-1] != new_unpad: - import cv2 - img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) - - # 添加边框 - top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) - left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) - img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) - - return img, r, (dw, dh) - - def get_class_names(self): - """获取类别名称列表""" - if hasattr(self, 'using_onnx') and self.using_onnx: - return self.class_names - elif hasattr(self, 'using_fallback') and self.using_fallback: - return self.names - elif hasattr(self, 'names'): - return self.names - else: - return self.model.names - def get_performance_stats(self): """获取性能统计信息""" if not self.inference_times: - return {"average_time": 0, "min_time": 0, "max_time": 0, "count": 0, "fps": 0} + return { + "average_time": 0, + "min_time": 0, + "max_time": 0, + "fps": 0 + } avg_time = sum(self.inference_times) / len(self.inference_times) - - # 修复fps计算逻辑 - if avg_time > 0: - fps = 1.0 / avg_time - else: - fps = 0 - return { "average_time": avg_time, "min_time": min(self.inference_times), "max_time": max(self.inference_times), - "count": len(self.inference_times), - "fps": fps - } \ No newline at end of file + "median_time": statistics.median(self.inference_times), + "fps": 1.0 / avg_time if avg_time > 0 else 0 + } + +if __name__ == "__main__": + # 简单测试 + logging.basicConfig(level=logging.INFO) + + try: + # 初始化模型 + model = YOLOModel() + + # 加载测试图像 + test_image_path = "data/training/images/test/sample_1.jpg" + if os.path.exists(test_image_path): + image = Image.open(test_image_path) + + # 检测 + detections = model.detect(image) + + # 打印结果 + print(f"检测到 {len(detections)} 个目标:") + for det in detections: + print(f"类别: {det['class_name']}, 置信度: {det['confidence']:.2f}, 边界框: {det['bbox']}") + + # 打印性能 + perf = model.get_performance_stats() + print(f"模型性能: {perf['average_time']*1000:.2f}ms, {perf['fps']:.2f} FPS") + else: + print(f"测试图像不存在: {test_image_path}") + + except Exception as e: + print(f"测试失败: {e}") \ No newline at end of file