Add NVIDIA GPU detection and TensorRT engine model support

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
+*.engine filter=lfs diff=lfs merge=lfs -text

check_model.py

@@ -15,7 +15,10 @@ logger = logging.getLogger(__name__)
 
 def check_model():
     """检查模型文件"""
-    model_path = "models/kunin-mice-pose.v0.1.5n.pt"
+    base_model_path = "models/kunin-mice-pose.v0.1.5n.pt"
+    
+    # 选择模型路径（与SingleMouseProcessor保持一致）
+    model_path = select_model_path(base_model_path)
     
     logger.info(f"检查模型文件: {model_path}")
     
@@ -87,6 +90,40 @@ def check_model():
         logger.error(traceback.format_exc())
         return False
 
+def select_model_path(base_model_path: str) -> str:
+    """根据GPU情况选择模型路径"""
+    try:
+        import torch
+        # 检测是否有NVIDIA GPU
+        if torch.cuda.is_available():
+            nvidia_gpu_found = False
+            for i in range(torch.cuda.device_count()):
+                gpu_name = torch.cuda.get_device_name(i).lower()
+                if 'nvidia' in gpu_name:
+                    nvidia_gpu_found = True
+                    break
+            
+            if nvidia_gpu_found:
+                # 构建.engine模型路径
+                engine_path = base_model_path.replace('.pt', '.engine')
+                if os.path.exists(engine_path):
+                    logger.info(f"🚀 检测到NVIDIA GPU，使用TensorRT模型: {engine_path}")
+                    return engine_path
+                else:
+                    logger.info(f"⚠️ NVIDIA GPU已检测到，但TensorRT模型不存在: {engine_path}")
+                    logger.info(f"📍 使用PyTorch模型: {base_model_path}")
+                    return base_model_path
+            else:
+                logger.info(f"📍 检测到GPU但非NVIDIA，使用PyTorch模型: {base_model_path}")
+                return base_model_path
+        else:
+            logger.info(f"📍 未检测到GPU，使用CPU模式，PyTorch模型: {base_model_path}")
+            return base_model_path
+            
+    except Exception as e:
+        logger.warning(f"⚠️ GPU检测失败，使用默认模型: {str(e)}")
+        return base_model_path
+
 def main():
     """主函数"""
     logger.info("🔍 开始模型检查...")

gradio_webrtc_server.py

@@ -24,7 +24,8 @@ class SingleMouseProcessor:
     """单鼠姿态检测处理器"""
     
     def __init__(self, model_path: str = "models/kunin-mice-pose.v0.1.5n.pt"):
-        self.model_path = model_path
+        self.base_model_path = model_path
+        self.model_path = self._select_model_path()
         self.model = None
         self.lock = Lock()
         self.frame_count = 0
@@ -57,6 +58,40 @@ class SingleMouseProcessor:
         # 加载模型
         self._load_model()
 
+    def _select_model_path(self) -> str:
+        """根据GPU情况选择模型路径"""
+        try:
+            import torch
+            # 检测是否有NVIDIA GPU
+            if torch.cuda.is_available():
+                nvidia_gpu_found = False
+                for i in range(torch.cuda.device_count()):
+                    gpu_name = torch.cuda.get_device_name(i).lower()
+                    if 'nvidia' in gpu_name:
+                        nvidia_gpu_found = True
+                        break
+                
+                if nvidia_gpu_found:
+                    # 构建.engine模型路径
+                    engine_path = self.base_model_path.replace('.pt', '.engine')
+                    if os.path.exists(engine_path):
+                        logger.info(f"🚀 检测到NVIDIA GPU，使用TensorRT模型: {engine_path}")
+                        return engine_path
+                    else:
+                        logger.info(f"⚠️ NVIDIA GPU已检测到，但TensorRT模型不存在: {engine_path}")
+                        logger.info(f"📍 使用PyTorch模型: {self.base_model_path}")
+                        return self.base_model_path
+                else:
+                    logger.info(f"📍 检测到GPU但非NVIDIA，使用PyTorch模型: {self.base_model_path}")
+                    return self.base_model_path
+            else:
+                logger.info(f"📍 未检测到GPU，使用CPU模式，PyTorch模型: {self.base_model_path}")
+                return self.base_model_path
+                
+        except Exception as e:
+            logger.warning(f"⚠️ GPU检测失败，使用默认模型: {str(e)}")
+            return self.base_model_path
+
     def _load_model(self):
         """加载YOLO模型"""
         try:

inspect_model.py

@@ -15,7 +15,10 @@ logger = logging.getLogger(__name__)
 
 def inspect_model():
     """检查模型详细信息"""
-    model_path = "models/kunin-mice-pose.v0.1.5n.pt"
+    base_model_path = "models/kunin-mice-pose.v0.1.5n.pt"
+    
+    # 选择模型路径（与SingleMouseProcessor保持一致）
+    model_path = select_model_path(base_model_path)
     
     try:
         from ultralytics import YOLO
@@ -122,6 +125,40 @@ def inspect_model():
         logger.error(traceback.format_exc())
         return False
 
+def select_model_path(base_model_path: str) -> str:
+    """根据GPU情况选择模型路径"""
+    try:
+        import torch
+        # 检测是否有NVIDIA GPU
+        if torch.cuda.is_available():
+            nvidia_gpu_found = False
+            for i in range(torch.cuda.device_count()):
+                gpu_name = torch.cuda.get_device_name(i).lower()
+                if 'nvidia' in gpu_name:
+                    nvidia_gpu_found = True
+                    break
+            
+            if nvidia_gpu_found:
+                # 构建.engine模型路径
+                engine_path = base_model_path.replace('.pt', '.engine')
+                if os.path.exists(engine_path):
+                    logger.info(f"🚀 检测到NVIDIA GPU，使用TensorRT模型: {engine_path}")
+                    return engine_path
+                else:
+                    logger.info(f"⚠️ NVIDIA GPU已检测到，但TensorRT模型不存在: {engine_path}")
+                    logger.info(f"📍 使用PyTorch模型: {base_model_path}")
+                    return base_model_path
+            else:
+                logger.info(f"📍 检测到GPU但非NVIDIA，使用PyTorch模型: {base_model_path}")
+                return base_model_path
+        else:
+            logger.info(f"📍 未检测到GPU，使用CPU模式，PyTorch模型: {base_model_path}")
+            return base_model_path
+            
+    except Exception as e:
+        logger.warning(f"⚠️ GPU检测失败，使用默认模型: {str(e)}")
+        return base_model_path
+
 def main():
     """主函数"""
     logger.info("🔍 开始模型详细检查...")