├── requirements.txt ├── .gitignore ├── mcp-config.json ├── LICENSE ├── setup.py ├── Readme.md ├── server_combine_terminal.py ├── server.py └── server_cli.py /requirements.txt: -------------------------------------------------------------------------------- 1 | mcp 2 | ultralytics 3 | opencv-python 4 | numpy 5 | pillow -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Python-generated files 2 | __pycache__/ 3 | *.py[oc] 4 | build/ 5 | dist/ 6 | wheels/ 7 | *.egg-info 8 | 9 | # Virtual environments 10 | .venv 11 | yolo_service.log 12 | *.pt 13 | runs/ -------------------------------------------------------------------------------- /mcp-config.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "yolo-service": { 4 | "command": "D:\\BackDataService\\YOLO-MCP-Server\\.venv\\Scripts\\python.exe", 5 | "args": [ 6 | "D:\\BackDataService\\YOLO-MCP-Server\\server.py" 7 | ], 8 | "env": { 9 | "PYTHONPATH": "D:\\BackDataService\\YOLO-MCP-Server" 10 | } 11 | } 12 | } 13 | } 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 GongRzhe 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Import necessary Python standard libraries 2 | import os # For operating with file system, handling files and directory paths 3 | import json # For processing JSON format data 4 | import subprocess # For creating and managing subprocesses 5 | import sys # For accessing Python interpreter related variables and functions 6 | import platform # For getting current operating system information 7 | 8 | def setup_venv(): 9 | """ 10 | Function to set up Python virtual environment 11 | 12 | Features: 13 | - Checks if Python version meets requirements (3.10+) 14 | - Creates Python virtual environment (if it doesn't exist) 15 | - Installs required dependencies in the newly created virtual environment 16 | 17 | No parameters required 18 | 19 | Returns: Path to Python interpreter in the virtual environment 20 | """ 21 | # Check Python version 22 | python_version = sys.version_info 23 | if python_version.major < 3 or (python_version.major == 3 and python_version.minor < 10): 24 | print("Error: Python 3.10 or higher is required.") 25 | sys.exit(1) 26 | 27 | # Get absolute path of the directory containing the current script 28 | base_path = os.path.abspath(os.path.dirname(__file__)) 29 | # Set virtual environment directory path, will create a directory named '.venv' under base_path 30 | venv_path = os.path.join(base_path, '.venv') 31 | # Flag whether a new virtual environment was created 32 | venv_created = False 33 | 34 | # Check if virtual environment already exists 35 | if not os.path.exists(venv_path): 36 | print("Creating virtual environment...") 37 | # Use Python's venv module to create virtual environment 38 | # sys.executable gets the path of the current Python interpreter 39 | subprocess.run([sys.executable, '-m', 'venv', venv_path], check=True) 40 | print("Virtual environment created successfully!") 41 | venv_created = True 42 | else: 43 | print("Virtual environment already exists.") 44 | 45 | # Determine pip and python executable paths based on operating system 46 | is_windows = platform.system() == "Windows" 47 | if is_windows: 48 | pip_path = os.path.join(venv_path, 'Scripts', 'pip.exe') 49 | python_path = os.path.join(venv_path, 'Scripts', 'python.exe') 50 | else: 51 | pip_path = os.path.join(venv_path, 'bin', 'pip') 52 | python_path = os.path.join(venv_path, 'bin', 'python') 53 | 54 | # Install or update dependencies 55 | print("\nInstalling requirements...") 56 | 57 | # Create requirements.txt with necessary packages for YOLO MCP server 58 | requirements = [ 59 | "mcp", # Model Context Protocol for server 60 | "ultralytics", # YOLO models 61 | "opencv-python", # For camera operations 62 | "numpy", # Numerical operations 63 | "pillow" # Image processing 64 | ] 65 | 66 | requirements_path = os.path.join(base_path, 'requirements.txt') 67 | with open(requirements_path, 'w') as f: 68 | f.write('\n'.join(requirements)) 69 | 70 | # Update pip using the python executable (more reliable method) 71 | try: 72 | subprocess.run([python_path, '-m', 'pip', 'install', '--upgrade', 'pip'], check=True) 73 | print("Pip upgraded successfully.") 74 | except subprocess.CalledProcessError: 75 | print("Warning: Pip upgrade failed, continuing with existing version.") 76 | 77 | # Install requirements 78 | subprocess.run([pip_path, 'install', '-r', requirements_path], check=True) 79 | 80 | print("Requirements installed successfully!") 81 | 82 | return python_path 83 | 84 | def generate_mcp_config(python_path): 85 | """ 86 | Function to generate MCP (Model Context Protocol) configuration file for YOLO service 87 | 88 | Features: 89 | - Creates configuration containing Python interpreter path and server script path 90 | - Saves configuration as JSON format file 91 | - Prints configuration information for different MCP clients 92 | 93 | Parameters: 94 | - python_path: Path to Python interpreter in the virtual environment 95 | 96 | Returns: None 97 | """ 98 | # Get absolute path of the directory containing the current script 99 | base_path = os.path.abspath(os.path.dirname(__file__)) 100 | 101 | # Path to YOLO MCP server script 102 | server_script_path = os.path.join(base_path, 'server.py') 103 | 104 | # Create MCP configuration dictionary 105 | config = { 106 | "mcpServers": { 107 | "yolo-service": { 108 | "command": python_path, 109 | "args": [server_script_path], 110 | "env": { 111 | "PYTHONPATH": base_path 112 | } 113 | } 114 | } 115 | } 116 | 117 | # Save configuration to JSON file 118 | config_path = os.path.join(base_path, 'mcp-config.json') 119 | with open(config_path, 'w') as f: 120 | json.dump(config, f, indent=2) # indent=2 gives the JSON file good formatting 121 | 122 | # Print configuration information 123 | print(f"\nMCP configuration has been written to: {config_path}") 124 | print(f"\nMCP configuration for Cursor:\n\n{python_path} {server_script_path}") 125 | print("\nMCP configuration for Windsurf/Claude Desktop:") 126 | print(json.dumps(config, indent=2)) 127 | 128 | # Provide instructions for adding configuration to Claude Desktop configuration file 129 | if platform.system() == "Windows": 130 | claude_config_path = os.path.expandvars("%APPDATA%\\Claude\\claude_desktop_config.json") 131 | else: # macOS 132 | claude_config_path = os.path.expanduser("~/Library/Application Support/Claude/claude_desktop_config.json") 133 | 134 | print(f"\nTo use with Claude Desktop, merge this configuration into: {claude_config_path}") 135 | 136 | # Code executed when the script is run directly (not imported) 137 | if __name__ == '__main__': 138 | # Execute main functions in sequence: 139 | # 1. Set up virtual environment and install dependencies 140 | python_path = setup_venv() 141 | # 2. Generate MCP configuration file 142 | generate_mcp_config(python_path) 143 | 144 | print("\nSetup complete! You can now use the YOLO MCP server with compatible clients.") -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 | # YOLO MCP Service 2 | 3 | A powerful YOLO (You Only Look Once) computer vision service that integrates with Claude AI through Model Context Protocol (MCP). This service enables Claude to perform object detection, segmentation, classification, and real-time camera analysis using state-of-the-art YOLO models. 4 | 5 | ![](https://badge.mcpx.dev?type=server 'MCP Server') 6 | 7 | 8 | ## Features 9 | 10 | - Object detection, segmentation, classification, and pose estimation 11 | - Real-time camera integration for live object detection 12 | - Support for model training, validation, and export 13 | - Comprehensive image analysis combining multiple models 14 | - Support for both file paths and base64-encoded images 15 | - Seamless integration with Claude AI 16 | 17 | ## Setup Instructions 18 | 19 | ### Prerequisites 20 | 21 | - Python 3.10 or higher 22 | - Git (optional, for cloning the repository) 23 | 24 | ### Environment Setup 25 | 26 | 1. Create a directory for the project and navigate to it: 27 | ```bash 28 | mkdir yolo-mcp-service 29 | cd yolo-mcp-service 30 | ``` 31 | 32 | 2. Download the project files or clone from repository: 33 | ```bash 34 | # If you have the files, copy them to this directory 35 | # If using git: 36 | git clone https://github.com/GongRzhe/YOLO-MCP-Server.git . 37 | ``` 38 | 39 | 3. Create a virtual environment: 40 | ```bash 41 | # On Windows 42 | python -m venv .venv 43 | 44 | # On macOS/Linux 45 | python3 -m venv .venv 46 | ``` 47 | 48 | 4. Activate the virtual environment: 49 | ```bash 50 | # On Windows 51 | .venv\Scripts\activate 52 | 53 | # On macOS/Linux 54 | source .venv/bin/activate 55 | ``` 56 | 57 | 5. Run the setup script: 58 | ```bash 59 | python setup.py 60 | ``` 61 | 62 | The setup script will: 63 | - Check your Python version 64 | - Create a virtual environment (if not already created) 65 | - Install required dependencies 66 | - Generate an MCP configuration file (mcp-config.json) 67 | - Output configuration information for different MCP clients including Claude 68 | 69 | 6. Note the output from the setup script, which will look similar to: 70 | ``` 71 | MCP configuration has been written to: /path/to/mcp-config.json 72 | 73 | MCP configuration for Cursor: 74 | 75 | /path/to/.venv/bin/python /path/to/server.py 76 | 77 | MCP configuration for Windsurf/Claude Desktop: 78 | { 79 | "mcpServers": { 80 | "yolo-service": { 81 | "command": "/path/to/.venv/bin/python", 82 | "args": [ 83 | "/path/to/server.py" 84 | ], 85 | "env": { 86 | "PYTHONPATH": "/path/to" 87 | } 88 | } 89 | } 90 | } 91 | 92 | To use with Claude Desktop, merge this configuration into: /path/to/claude_desktop_config.json 93 | ``` 94 | 95 | ### Downloading YOLO Models 96 | 97 | Before using the service, you need to download the YOLO models. The service looks for models in the following directories: 98 | - The current directory where the service is running 99 | - A `models` subdirectory 100 | - Any other directory configured in the `CONFIG["model_dirs"]` variable in server.py 101 | 102 | Create a models directory and download some common models: 103 | 104 | ```bash 105 | # Create models directory 106 | mkdir models 107 | 108 | # Download YOLOv8n for basic object detection 109 | curl -L https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt -o models/yolov8n.pt 110 | 111 | # Download YOLOv8n-seg for segmentation 112 | curl -L https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt -o models/yolov8n-seg.pt 113 | 114 | # Download YOLOv8n-cls for classification 115 | curl -L https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt -o models/yolov8n-cls.pt 116 | 117 | # Download YOLOv8n-pose for pose estimation 118 | curl -L https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-pose.pt -o models/yolov8n-pose.pt 119 | ``` 120 | 121 | For Windows PowerShell users: 122 | ```powershell 123 | # Create models directory 124 | mkdir models 125 | 126 | # Download models using Invoke-WebRequest 127 | Invoke-WebRequest -Uri "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt" -OutFile "models/yolov8n.pt" 128 | Invoke-WebRequest -Uri "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt" -OutFile "models/yolov8n-seg.pt" 129 | Invoke-WebRequest -Uri "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt" -OutFile "models/yolov8n-cls.pt" 130 | Invoke-WebRequest -Uri "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-pose.pt" -OutFile "models/yolov8n-pose.pt" 131 | ``` 132 | 133 | ### Configuring Claude 134 | 135 | To use this service with Claude: 136 | 137 | 1. For Claude web: Set up the service on your local machine and use the configuration provided by the setup script in your MCP client. 138 | 139 | 2. For Claude Desktop: 140 | - Run the setup script and note the configuration output 141 | - Locate your Claude Desktop configuration file (the path is provided in the setup script output) 142 | - Add or merge the configuration into your Claude Desktop configuration file 143 | - Restart Claude Desktop 144 | 145 | ## Using YOLO Tools in Claude 146 | 147 | ### 1. First Check Available Models 148 | 149 | Always check which models are available on your system first: 150 | 151 | ``` 152 | I'd like to use the YOLO tools. Can you first check which models are available on my system? 153 | 154 | 155 | 156 | 157 | 158 | ``` 159 | 160 | ### 2. Detecting Objects in an Image 161 | 162 | For analyzing an image file on your computer: 163 | 164 | ``` 165 | Can you analyze this image file for objects? 166 | 167 | 168 | 169 | /path/to/your/image.jpg 170 | 0.3 171 | 172 | 173 | ``` 174 | 175 | You can also specify a different model: 176 | 177 | ``` 178 | Can you analyze this image using a different model? 179 | 180 | 181 | 182 | /path/to/your/image.jpg 183 | yolov8n.pt 184 | 0.4 185 | 186 | 187 | ``` 188 | 189 | ### 3. Running Comprehensive Image Analysis 190 | 191 | For more detailed analysis that combines object detection, classification, and more: 192 | 193 | ``` 194 | Can you perform a comprehensive analysis on this image? 195 | 196 | 197 | 198 | /path/to/your/image.jpg 199 | 0.3 200 | 201 | 202 | ``` 203 | 204 | ### 4. Image Segmentation 205 | 206 | For identifying object boundaries and creating segmentation masks: 207 | 208 | ``` 209 | Can you perform image segmentation on this photo? 210 | 211 | 212 | 213 | /path/to/your/image.jpg 214 | true 215 | yolov8n-seg.pt 216 | 217 | 218 | ``` 219 | 220 | ### 5. Image Classification 221 | 222 | For classifying the entire image content: 223 | 224 | ``` 225 | What does this image show? Can you classify it? 226 | 227 | 228 | 229 | /path/to/your/image.jpg 230 | true 231 | yolov8n-cls.pt 232 | 5 233 | 234 | 235 | ``` 236 | 237 | ### 6. Using Your Computer's Camera 238 | 239 | Start real-time object detection using your computer's camera: 240 | 241 | ``` 242 | Can you turn on my camera and detect objects in real-time? 243 | 244 | 245 | 246 | yolov8n.pt 247 | 0.3 248 | 249 | 250 | ``` 251 | 252 | Get the latest camera detections: 253 | 254 | ``` 255 | What are you seeing through my camera right now? 256 | 257 | 258 | 259 | 260 | 261 | ``` 262 | 263 | Stop the camera when finished: 264 | 265 | ``` 266 | Please turn off the camera. 267 | 268 | 269 | 270 | 271 | 272 | ``` 273 | 274 | ### 7. Advanced Model Operations 275 | 276 | #### Training a Custom Model 277 | 278 | ``` 279 | I want to train a custom object detection model on my dataset. 280 | 281 | 282 | 283 | /path/to/your/dataset 284 | yolov8n.pt 285 | 50 286 | 287 | 288 | ``` 289 | 290 | #### Validating a Model 291 | 292 | ``` 293 | Can you validate the performance of my model on a test dataset? 294 | 295 | 296 | 297 | /path/to/your/trained/model.pt 298 | /path/to/validation/dataset 299 | 300 | 301 | ``` 302 | 303 | #### Exporting a Model to Different Formats 304 | 305 | ``` 306 | I need to export my YOLO model to ONNX format. 307 | 308 | 309 | 310 | /path/to/your/model.pt 311 | onnx 312 | 313 | 314 | ``` 315 | 316 | ### 8. Testing Connection 317 | 318 | Check if the YOLO service is running correctly: 319 | 320 | ``` 321 | Is the YOLO service running correctly? 322 | 323 | 324 | 325 | 326 | 327 | ``` 328 | 329 | ## Troubleshooting 330 | 331 | ### Camera Issues 332 | 333 | If the camera doesn't work, try different camera IDs: 334 | 335 | ``` 336 | 337 | 338 | 1 339 | 340 | 341 | ``` 342 | 343 | ### Model Not Found 344 | 345 | If a model is not found, make sure you've downloaded it to one of the configured directories: 346 | 347 | ``` 348 | 349 | 350 | 351 | 352 | ``` 353 | 354 | ### Performance Issues 355 | 356 | For better performance with limited resources, use the smaller models (e.g., yolov8n.pt instead of yolov8x.pt) 357 | -------------------------------------------------------------------------------- /server_combine_terminal.py: -------------------------------------------------------------------------------- 1 | # server.py - CLI version (command return only) 2 | import fnmatch 3 | import os 4 | import base64 5 | import time 6 | import threading 7 | import json 8 | import tempfile 9 | import platform 10 | from io import BytesIO 11 | from typing import List, Dict, Any, Optional, Union 12 | import numpy as np 13 | from PIL import Image 14 | 15 | from mcp.server.fastmcp import FastMCP 16 | 17 | # Set up logging configuration 18 | import os.path 19 | import sys 20 | import logging 21 | import contextlib 22 | import signal 23 | import atexit 24 | 25 | logging.basicConfig( 26 | level=logging.INFO, 27 | format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', 28 | handlers=[ 29 | logging.FileHandler("yolo_service.log"), 30 | logging.StreamHandler(sys.stderr) 31 | ] 32 | ) 33 | camera_startup_status = None # Will store error details if startup fails 34 | camera_last_error = None 35 | logger = logging.getLogger('yolo_service') 36 | 37 | # Global variables for camera control 38 | camera_running = False 39 | camera_thread = None 40 | detection_results = [] 41 | camera_last_access_time = 0 42 | CAMERA_INACTIVITY_TIMEOUT = 60 # Auto-shutdown after 60 seconds of inactivity 43 | 44 | def load_image(image_source, is_path=False): 45 | """ 46 | Load image from file path or base64 data 47 | 48 | Args: 49 | image_source: File path or base64 encoded image data 50 | is_path: Whether image_source is a file path 51 | 52 | Returns: 53 | PIL Image object 54 | """ 55 | try: 56 | if is_path: 57 | # Load image from file path 58 | if os.path.exists(image_source): 59 | return Image.open(image_source) 60 | else: 61 | raise FileNotFoundError(f"Image file not found: {image_source}") 62 | else: 63 | # Load image from base64 data 64 | image_bytes = base64.b64decode(image_source) 65 | return Image.open(BytesIO(image_bytes)) 66 | except Exception as e: 67 | raise ValueError(f"Failed to load image: {str(e)}") 68 | 69 | # Modified function to just return the command string 70 | def run_yolo_cli(command_args, capture_output=True, timeout=60): 71 | """ 72 | Return the YOLO CLI command string without executing it 73 | 74 | Args: 75 | command_args: List of command arguments to pass to yolo CLI 76 | capture_output: Not used, kept for compatibility with original function 77 | timeout: Not used, kept for compatibility with original function 78 | 79 | Returns: 80 | Dictionary containing the command string 81 | """ 82 | # Build the complete command 83 | cmd = ["yolo"] + command_args 84 | cmd_str = " ".join(cmd) 85 | 86 | # Log the command 87 | logger.info(f"Would run YOLO CLI command: {cmd_str}") 88 | 89 | # Return the command string in a similar structure as the original function 90 | return { 91 | "success": True, 92 | "command": cmd_str, 93 | "would_execute": True, 94 | "note": "CLI execution disabled, showing command only" 95 | } 96 | 97 | # Create MCP server 98 | mcp = FastMCP("YOLO_Service") 99 | 100 | # Global configuration 101 | CONFIG = { 102 | "model_dirs": [ 103 | ".", # Current directory 104 | "./models", # Models subdirectory 105 | os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"), 106 | ] 107 | } 108 | 109 | # Function to save base64 data to temp file 110 | def save_base64_to_temp(base64_data, prefix="image", suffix=".jpg"): 111 | """Save base64 encoded data to a temporary file and return the path""" 112 | try: 113 | # Create a temporary file 114 | fd, temp_path = tempfile.mkstemp(suffix=suffix, prefix=prefix) 115 | 116 | # Decode base64 data 117 | image_data = base64.b64decode(base64_data) 118 | 119 | # Write data to file 120 | with os.fdopen(fd, 'wb') as temp_file: 121 | temp_file.write(image_data) 122 | 123 | return temp_path 124 | except Exception as e: 125 | logger.error(f"Error saving base64 to temp file: {str(e)}") 126 | raise ValueError(f"Failed to save base64 data: {str(e)}") 127 | 128 | @mcp.tool() 129 | def get_model_directories() -> Dict[str, Any]: 130 | """Get information about configured model directories and available models""" 131 | directories = [] 132 | 133 | for directory in CONFIG["model_dirs"]: 134 | dir_info = { 135 | "path": directory, 136 | "exists": os.path.exists(directory), 137 | "is_directory": os.path.isdir(directory) if os.path.exists(directory) else False, 138 | "models": [] 139 | } 140 | 141 | if dir_info["exists"] and dir_info["is_directory"]: 142 | for filename in os.listdir(directory): 143 | if filename.endswith(".pt"): 144 | dir_info["models"].append(filename) 145 | 146 | directories.append(dir_info) 147 | 148 | return { 149 | "configured_directories": CONFIG["model_dirs"], 150 | "directory_details": directories, 151 | "available_models": list_available_models(), 152 | "loaded_models": [] # No longer track loaded models with CLI approach 153 | } 154 | 155 | @mcp.tool() 156 | def detect_objects( 157 | image_data: str, 158 | model_name: str = "yolov8n.pt", 159 | confidence: float = 0.25, 160 | save_results: bool = False, 161 | is_path: bool = False 162 | ) -> Dict[str, Any]: 163 | """ 164 | Return the YOLO CLI command for object detection without executing it 165 | 166 | Args: 167 | image_data: Base64 encoded image or file path (if is_path=True) 168 | model_name: YOLO model name 169 | confidence: Detection confidence threshold 170 | save_results: Whether to save results to disk 171 | is_path: Whether image_data is a file path 172 | 173 | Returns: 174 | Dictionary containing command that would be executed 175 | """ 176 | try: 177 | # Determine source path 178 | if is_path: 179 | source_path = image_data 180 | if not os.path.exists(source_path): 181 | return { 182 | "error": f"Image file not found: {source_path}", 183 | "source": source_path 184 | } 185 | else: 186 | # For base64, we would save to temp file, but we'll just indicate this 187 | source_path = "[temp_file_from_base64]" 188 | 189 | # Determine full model path 190 | model_path = None 191 | for directory in CONFIG["model_dirs"]: 192 | potential_path = os.path.join(directory, model_name) 193 | if os.path.exists(potential_path): 194 | model_path = potential_path 195 | break 196 | 197 | if model_path is None: 198 | available = list_available_models() 199 | available_str = ", ".join(available) if available else "none" 200 | return { 201 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}", 202 | "source": image_data if is_path else "base64_image" 203 | } 204 | 205 | # Setup output directory for save_results 206 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_results") 207 | 208 | # Build YOLO CLI command 209 | cmd_args = [ 210 | "detect", # Task 211 | "predict", # Mode 212 | f"model={model_path}", 213 | f"source={source_path}", 214 | f"conf={confidence}", 215 | "format=json", # Request JSON output for parsing 216 | ] 217 | 218 | if save_results: 219 | cmd_args.append(f"project={output_dir}") 220 | cmd_args.append("save=True") 221 | else: 222 | cmd_args.append("save=False") 223 | 224 | # Get command string without executing 225 | result = run_yolo_cli(cmd_args) 226 | 227 | # Return command information 228 | return { 229 | "status": "command_generated", 230 | "model_used": model_name, 231 | "model_path": model_path, 232 | "source": source_path, 233 | "command": result["command"], 234 | "note": "Command generated but not executed - detection results would be returned from actual execution", 235 | "parameters": { 236 | "confidence": confidence, 237 | "save_results": save_results, 238 | "is_path": is_path, 239 | "output_dir": output_dir if save_results else None 240 | } 241 | } 242 | 243 | except Exception as e: 244 | logger.error(f"Error in detect_objects command generation: {str(e)}") 245 | return { 246 | "error": f"Failed to generate detection command: {str(e)}", 247 | "source": image_data if is_path else "base64_image" 248 | } 249 | 250 | @mcp.tool() 251 | def segment_objects( 252 | image_data: str, 253 | model_name: str = "yolov11n-seg.pt", 254 | confidence: float = 0.25, 255 | save_results: bool = False, 256 | is_path: bool = False 257 | ) -> Dict[str, Any]: 258 | """ 259 | Return the YOLO CLI command for segmentation without executing it 260 | 261 | Args: 262 | image_data: Base64 encoded image or file path (if is_path=True) 263 | model_name: YOLO segmentation model name 264 | confidence: Detection confidence threshold 265 | save_results: Whether to save results to disk 266 | is_path: Whether image_data is a file path 267 | 268 | Returns: 269 | Dictionary containing command that would be executed 270 | """ 271 | try: 272 | # Determine source path 273 | if is_path: 274 | source_path = image_data 275 | if not os.path.exists(source_path): 276 | return { 277 | "error": f"Image file not found: {source_path}", 278 | "source": source_path 279 | } 280 | else: 281 | # For base64, we would save to temp file, but we'll just indicate this 282 | source_path = "[temp_file_from_base64]" 283 | 284 | # Determine full model path 285 | model_path = None 286 | for directory in CONFIG["model_dirs"]: 287 | potential_path = os.path.join(directory, model_name) 288 | if os.path.exists(potential_path): 289 | model_path = potential_path 290 | break 291 | 292 | if model_path is None: 293 | available = list_available_models() 294 | available_str = ", ".join(available) if available else "none" 295 | return { 296 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}", 297 | "source": image_data if is_path else "base64_image" 298 | } 299 | 300 | # Setup output directory for save_results 301 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_results") 302 | 303 | # Build YOLO CLI command 304 | cmd_args = [ 305 | "segment", # Task 306 | "predict", # Mode 307 | f"model={model_path}", 308 | f"source={source_path}", 309 | f"conf={confidence}", 310 | "format=json", # Request JSON output for parsing 311 | ] 312 | 313 | if save_results: 314 | cmd_args.append(f"project={output_dir}") 315 | cmd_args.append("save=True") 316 | else: 317 | cmd_args.append("save=False") 318 | 319 | # Get command string without executing 320 | result = run_yolo_cli(cmd_args) 321 | 322 | # Return command information 323 | return { 324 | "status": "command_generated", 325 | "model_used": model_name, 326 | "model_path": model_path, 327 | "source": source_path, 328 | "command": result["command"], 329 | "note": "Command generated but not executed - segmentation results would be returned from actual execution", 330 | "parameters": { 331 | "confidence": confidence, 332 | "save_results": save_results, 333 | "is_path": is_path, 334 | "output_dir": output_dir if save_results else None 335 | } 336 | } 337 | 338 | except Exception as e: 339 | logger.error(f"Error in segment_objects command generation: {str(e)}") 340 | return { 341 | "error": f"Failed to generate segmentation command: {str(e)}", 342 | "source": image_data if is_path else "base64_image" 343 | } 344 | 345 | @mcp.tool() 346 | def classify_image( 347 | image_data: str, 348 | model_name: str = "yolov11n-cls.pt", 349 | top_k: int = 5, 350 | save_results: bool = False, 351 | is_path: bool = False 352 | ) -> Dict[str, Any]: 353 | """ 354 | Return the YOLO CLI command for image classification without executing it 355 | 356 | Args: 357 | image_data: Base64 encoded image or file path (if is_path=True) 358 | model_name: YOLO classification model name 359 | top_k: Number of top categories to return 360 | save_results: Whether to save results to disk 361 | is_path: Whether image_data is a file path 362 | 363 | Returns: 364 | Dictionary containing command that would be executed 365 | """ 366 | try: 367 | # Determine source path 368 | if is_path: 369 | source_path = image_data 370 | if not os.path.exists(source_path): 371 | return { 372 | "error": f"Image file not found: {source_path}", 373 | "source": source_path 374 | } 375 | else: 376 | # For base64, we would save to temp file, but we'll just indicate this 377 | source_path = "[temp_file_from_base64]" 378 | 379 | # Determine full model path 380 | model_path = None 381 | for directory in CONFIG["model_dirs"]: 382 | potential_path = os.path.join(directory, model_name) 383 | if os.path.exists(potential_path): 384 | model_path = potential_path 385 | break 386 | 387 | if model_path is None: 388 | available = list_available_models() 389 | available_str = ", ".join(available) if available else "none" 390 | return { 391 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}", 392 | "source": image_data if is_path else "base64_image" 393 | } 394 | 395 | # Setup output directory for save_results 396 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_results") 397 | 398 | # Build YOLO CLI command 399 | cmd_args = [ 400 | "classify", # Task 401 | "predict", # Mode 402 | f"model={model_path}", 403 | f"source={source_path}", 404 | "format=json", # Request JSON output for parsing 405 | ] 406 | 407 | if save_results: 408 | cmd_args.append(f"project={output_dir}") 409 | cmd_args.append("save=True") 410 | else: 411 | cmd_args.append("save=False") 412 | 413 | # Get command string without executing 414 | result = run_yolo_cli(cmd_args) 415 | 416 | # Return command information 417 | return { 418 | "status": "command_generated", 419 | "model_used": model_name, 420 | "model_path": model_path, 421 | "source": source_path, 422 | "command": result["command"], 423 | "note": "Command generated but not executed - classification results would be returned from actual execution", 424 | "parameters": { 425 | "top_k": top_k, 426 | "save_results": save_results, 427 | "is_path": is_path, 428 | "output_dir": output_dir if save_results else None 429 | } 430 | } 431 | 432 | except Exception as e: 433 | logger.error(f"Error in classify_image command generation: {str(e)}") 434 | return { 435 | "error": f"Failed to generate classification command: {str(e)}", 436 | "source": image_data if is_path else "base64_image" 437 | } 438 | 439 | @mcp.tool() 440 | def track_objects( 441 | image_data: str, 442 | model_name: str = "yolov8n.pt", 443 | confidence: float = 0.25, 444 | tracker: str = "bytetrack.yaml", 445 | save_results: bool = False 446 | ) -> Dict[str, Any]: 447 | """ 448 | Return the YOLO CLI command for object tracking without executing it 449 | 450 | Args: 451 | image_data: Base64 encoded image 452 | model_name: YOLO model name 453 | confidence: Detection confidence threshold 454 | tracker: Tracker name to use (e.g., 'bytetrack.yaml', 'botsort.yaml') 455 | save_results: Whether to save results to disk 456 | 457 | Returns: 458 | Dictionary containing command that would be executed 459 | """ 460 | try: 461 | # For base64, we would save to temp file, but we'll just indicate this 462 | source_path = "[temp_file_from_base64]" 463 | 464 | # Determine full model path 465 | model_path = None 466 | for directory in CONFIG["model_dirs"]: 467 | potential_path = os.path.join(directory, model_name) 468 | if os.path.exists(potential_path): 469 | model_path = potential_path 470 | break 471 | 472 | if model_path is None: 473 | available = list_available_models() 474 | available_str = ", ".join(available) if available else "none" 475 | return { 476 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}" 477 | } 478 | 479 | # Setup output directory for save_results 480 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_track_results") 481 | 482 | # Build YOLO CLI command 483 | cmd_args = [ 484 | "track", # Combined task and mode for tracking 485 | f"model={model_path}", 486 | f"source={source_path}", 487 | f"conf={confidence}", 488 | f"tracker={tracker}", 489 | "format=json", # Request JSON output for parsing 490 | ] 491 | 492 | if save_results: 493 | cmd_args.append(f"project={output_dir}") 494 | cmd_args.append("save=True") 495 | else: 496 | cmd_args.append("save=False") 497 | 498 | # Get command string without executing 499 | result = run_yolo_cli(cmd_args) 500 | 501 | # Return command information 502 | return { 503 | "status": "command_generated", 504 | "model_used": model_name, 505 | "model_path": model_path, 506 | "source": source_path, 507 | "command": result["command"], 508 | "note": "Command generated but not executed - tracking results would be returned from actual execution", 509 | "parameters": { 510 | "confidence": confidence, 511 | "tracker": tracker, 512 | "save_results": save_results, 513 | "output_dir": output_dir if save_results else None 514 | } 515 | } 516 | 517 | except Exception as e: 518 | logger.error(f"Error in track_objects command generation: {str(e)}") 519 | return { 520 | "error": f"Failed to generate tracking command: {str(e)}" 521 | } 522 | 523 | @mcp.tool() 524 | def train_model( 525 | dataset_path: str, 526 | model_name: str = "yolov8n.pt", 527 | epochs: int = 100, 528 | imgsz: int = 640, 529 | batch: int = 16, 530 | name: str = "yolo_custom_model", 531 | project: str = "runs/train" 532 | ) -> Dict[str, Any]: 533 | """ 534 | Return the YOLO CLI command for model training without executing it 535 | 536 | Args: 537 | dataset_path: Path to YOLO format dataset 538 | model_name: Base model to start with 539 | epochs: Number of training epochs 540 | imgsz: Image size for training 541 | batch: Batch size 542 | name: Name for the training run 543 | project: Project directory 544 | 545 | Returns: 546 | Dictionary containing command that would be executed 547 | """ 548 | # Validate dataset path 549 | if not os.path.exists(dataset_path): 550 | return {"error": f"Dataset not found: {dataset_path}"} 551 | 552 | # Determine full model path 553 | model_path = None 554 | for directory in CONFIG["model_dirs"]: 555 | potential_path = os.path.join(directory, model_name) 556 | if os.path.exists(potential_path): 557 | model_path = potential_path 558 | break 559 | 560 | if model_path is None: 561 | available = list_available_models() 562 | available_str = ", ".join(available) if available else "none" 563 | return { 564 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}" 565 | } 566 | 567 | # Determine task type based on model name 568 | task = "detect" # Default task 569 | if "seg" in model_name: 570 | task = "segment" 571 | elif "pose" in model_name: 572 | task = "pose" 573 | elif "cls" in model_name: 574 | task = "classify" 575 | elif "obb" in model_name: 576 | task = "obb" 577 | 578 | # Build YOLO CLI command 579 | cmd_args = [ 580 | task, # Task 581 | "train", # Mode 582 | f"model={model_path}", 583 | f"data={dataset_path}", 584 | f"epochs={epochs}", 585 | f"imgsz={imgsz}", 586 | f"batch={batch}", 587 | f"name={name}", 588 | f"project={project}" 589 | ] 590 | 591 | # Get command string without executing 592 | result = run_yolo_cli(cmd_args) 593 | 594 | # Return command information 595 | return { 596 | "status": "command_generated", 597 | "model_used": model_name, 598 | "model_path": model_path, 599 | "command": result["command"], 600 | "note": "Command generated but not executed - training would start with actual execution", 601 | "parameters": { 602 | "dataset_path": dataset_path, 603 | "epochs": epochs, 604 | "imgsz": imgsz, 605 | "batch": batch, 606 | "name": name, 607 | "project": project, 608 | "task": task 609 | } 610 | } 611 | 612 | @mcp.tool() 613 | def validate_model( 614 | model_path: str, 615 | data_path: str, 616 | imgsz: int = 640, 617 | batch: int = 16 618 | ) -> Dict[str, Any]: 619 | """ 620 | Return the YOLO CLI command for model validation without executing it 621 | 622 | Args: 623 | model_path: Path to YOLO model (.pt file) 624 | data_path: Path to YOLO format validation dataset 625 | imgsz: Image size for validation 626 | batch: Batch size 627 | 628 | Returns: 629 | Dictionary containing command that would be executed 630 | """ 631 | # Validate model path 632 | if not os.path.exists(model_path): 633 | return {"error": f"Model file not found: {model_path}"} 634 | 635 | # Validate dataset path 636 | if not os.path.exists(data_path): 637 | return {"error": f"Dataset not found: {data_path}"} 638 | 639 | # Determine task type based on model name 640 | model_name = os.path.basename(model_path) 641 | task = "detect" # Default task 642 | if "seg" in model_name: 643 | task = "segment" 644 | elif "pose" in model_name: 645 | task = "pose" 646 | elif "cls" in model_name: 647 | task = "classify" 648 | elif "obb" in model_name: 649 | task = "obb" 650 | 651 | # Build YOLO CLI command 652 | cmd_args = [ 653 | task, # Task 654 | "val", # Mode 655 | f"model={model_path}", 656 | f"data={data_path}", 657 | f"imgsz={imgsz}", 658 | f"batch={batch}" 659 | ] 660 | 661 | # Get command string without executing 662 | result = run_yolo_cli(cmd_args) 663 | 664 | # Return command information 665 | return { 666 | "status": "command_generated", 667 | "model_path": model_path, 668 | "command": result["command"], 669 | "note": "Command generated but not executed - validation would begin with actual execution", 670 | "parameters": { 671 | "data_path": data_path, 672 | "imgsz": imgsz, 673 | "batch": batch, 674 | "task": task 675 | } 676 | } 677 | 678 | @mcp.tool() 679 | def export_model( 680 | model_path: str, 681 | format: str = "onnx", 682 | imgsz: int = 640 683 | ) -> Dict[str, Any]: 684 | """ 685 | Return the YOLO CLI command for model export without executing it 686 | 687 | Args: 688 | model_path: Path to YOLO model (.pt file) 689 | format: Export format (onnx, torchscript, openvino, etc.) 690 | imgsz: Image size for export 691 | 692 | Returns: 693 | Dictionary containing command that would be executed 694 | """ 695 | # Validate model path 696 | if not os.path.exists(model_path): 697 | return {"error": f"Model file not found: {model_path}"} 698 | 699 | # Valid export formats 700 | valid_formats = [ 701 | "torchscript", "onnx", "openvino", "engine", "coreml", "saved_model", 702 | "pb", "tflite", "edgetpu", "tfjs", "paddle" 703 | ] 704 | 705 | if format not in valid_formats: 706 | return {"error": f"Invalid export format: {format}. Valid formats include: {', '.join(valid_formats)}"} 707 | 708 | # Build YOLO CLI command 709 | cmd_args = [ 710 | "export", # Combined task and mode for export 711 | f"model={model_path}", 712 | f"format={format}", 713 | f"imgsz={imgsz}" 714 | ] 715 | 716 | # Get command string without executing 717 | result = run_yolo_cli(cmd_args) 718 | 719 | # Return command information 720 | return { 721 | "status": "command_generated", 722 | "model_path": model_path, 723 | "command": result["command"], 724 | "note": "Command generated but not executed - export would begin with actual execution", 725 | "parameters": { 726 | "format": format, 727 | "imgsz": imgsz, 728 | "expected_output": f"{os.path.splitext(model_path)[0]}.{format}" 729 | } 730 | } 731 | 732 | @mcp.tool() 733 | def list_available_models() -> List[str]: 734 | """List available YOLO models that actually exist on disk in any configured directory""" 735 | # Common YOLO model patterns 736 | model_patterns = [ 737 | "yolov11*.pt", 738 | "yolov8*.pt" 739 | ] 740 | 741 | # Find all existing models in all configured directories 742 | available_models = set() 743 | for directory in CONFIG["model_dirs"]: 744 | if not os.path.exists(directory): 745 | continue 746 | 747 | # Check for model files directly 748 | for filename in os.listdir(directory): 749 | if filename.endswith(".pt") and any( 750 | fnmatch.fnmatch(filename, pattern) for pattern in model_patterns 751 | ): 752 | available_models.add(filename) 753 | 754 | # Convert to sorted list 755 | result = sorted(list(available_models)) 756 | 757 | if not result: 758 | logger.warning("No model files found in configured directories.") 759 | return ["No models available - download models to any of these directories: " + ", ".join(CONFIG["model_dirs"])] 760 | 761 | return result 762 | 763 | @mcp.tool() 764 | def start_camera_detection( 765 | model_name: str = "yolov8n.pt", 766 | confidence: float = 0.25, 767 | camera_id: int = 0 768 | ) -> Dict[str, Any]: 769 | """ 770 | Return the YOLO CLI command for starting camera detection without executing it 771 | 772 | Args: 773 | model_name: YOLO model name to use 774 | confidence: Detection confidence threshold 775 | camera_id: Camera device ID (0 is usually the default camera) 776 | 777 | Returns: 778 | Dictionary containing command that would be executed 779 | """ 780 | # Determine full model path 781 | model_path = None 782 | for directory in CONFIG["model_dirs"]: 783 | potential_path = os.path.join(directory, model_name) 784 | if os.path.exists(potential_path): 785 | model_path = potential_path 786 | break 787 | 788 | if model_path is None: 789 | available = list_available_models() 790 | available_str = ", ".join(available) if available else "none" 791 | return { 792 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}" 793 | } 794 | 795 | # Determine task type based on model name 796 | task = "detect" # Default task 797 | if "seg" in model_name: 798 | task = "segment" 799 | elif "pose" in model_name: 800 | task = "pose" 801 | elif "cls" in model_name: 802 | task = "classify" 803 | 804 | # Build YOLO CLI command 805 | cmd_args = [ 806 | task, # Task 807 | "predict", # Mode 808 | f"model={model_path}", 809 | f"source={camera_id}", # Camera source ID 810 | f"conf={confidence}", 811 | "format=json", 812 | "save=False", # Don't save frames by default 813 | "show=True" # Show GUI window for camera view 814 | ] 815 | 816 | # Get command string without executing 817 | result = run_yolo_cli(cmd_args) 818 | 819 | # Return command information 820 | return { 821 | "status": "command_generated", 822 | "model_used": model_name, 823 | "model_path": model_path, 824 | "command": result["command"], 825 | "note": "Command generated but not executed - camera would start with actual execution", 826 | "parameters": { 827 | "confidence": confidence, 828 | "camera_id": camera_id, 829 | "task": task 830 | } 831 | } 832 | 833 | @mcp.tool() 834 | def stop_camera_detection() -> Dict[str, Any]: 835 | """ 836 | Simulate stopping camera detection (no actual command to execute) 837 | 838 | Returns: 839 | Information message 840 | """ 841 | return { 842 | "status": "command_generated", 843 | "message": "To stop camera detection, close the YOLO window or press 'q' in the terminal", 844 | "note": "Since commands are not executed, no actual camera is running" 845 | } 846 | 847 | @mcp.tool() 848 | def get_camera_detections() -> Dict[str, Any]: 849 | """ 850 | Simulate getting latest camera detections (no actual command to execute) 851 | 852 | Returns: 853 | Information message 854 | """ 855 | return { 856 | "status": "command_generated", 857 | "message": "Camera detections would be returned here if a camera was running", 858 | "note": "Since commands are not executed, no camera is running and no detections are available" 859 | } 860 | 861 | @mcp.tool() 862 | def comprehensive_image_analysis( 863 | image_path: str, 864 | confidence: float = 0.25, 865 | save_results: bool = False 866 | ) -> Dict[str, Any]: 867 | """ 868 | Return the YOLO CLI commands for comprehensive image analysis without executing them 869 | 870 | Args: 871 | image_path: Path to the image file 872 | confidence: Detection confidence threshold 873 | save_results: Whether to save results to disk 874 | 875 | Returns: 876 | Dictionary containing commands that would be executed 877 | """ 878 | if not os.path.exists(image_path): 879 | return {"error": f"Image file not found: {image_path}"} 880 | 881 | commands = [] 882 | 883 | # 1. Object detection 884 | detect_result = detect_objects( 885 | image_data=image_path, 886 | model_name="yolov11n.pt", 887 | confidence=confidence, 888 | save_results=save_results, 889 | is_path=True 890 | ) 891 | if "command" in detect_result: 892 | commands.append({ 893 | "task": "object_detection", 894 | "command": detect_result["command"] 895 | }) 896 | 897 | # 2. Scene classification 898 | try: 899 | cls_result = classify_image( 900 | image_data=image_path, 901 | model_name="yolov8n-cls.pt", 902 | top_k=3, 903 | save_results=save_results, 904 | is_path=True 905 | ) 906 | if "command" in cls_result: 907 | commands.append({ 908 | "task": "classification", 909 | "command": cls_result["command"] 910 | }) 911 | except Exception as e: 912 | logger.error(f"Error generating classification command: {str(e)}") 913 | 914 | # 3. Pose detection if available 915 | for directory in CONFIG["model_dirs"]: 916 | pose_model_path = os.path.join(directory, "yolov8n-pose.pt") 917 | if os.path.exists(pose_model_path): 918 | # Build YOLO CLI command for pose detection 919 | cmd_args = [ 920 | "pose", # Task 921 | "predict", # Mode 922 | f"model={pose_model_path}", 923 | f"source={image_path}", 924 | f"conf={confidence}", 925 | "format=json", 926 | ] 927 | 928 | if save_results: 929 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_pose_results") 930 | cmd_args.append(f"project={output_dir}") 931 | cmd_args.append("save=True") 932 | else: 933 | cmd_args.append("save=False") 934 | 935 | result = run_yolo_cli(cmd_args) 936 | 937 | commands.append({ 938 | "task": "pose_detection", 939 | "command": result["command"] 940 | }) 941 | break 942 | 943 | return { 944 | "status": "commands_generated", 945 | "image_path": image_path, 946 | "commands": commands, 947 | "note": "Commands generated but not executed - comprehensive analysis would occur with actual execution", 948 | "parameters": { 949 | "confidence": confidence, 950 | "save_results": save_results 951 | } 952 | } 953 | 954 | @mcp.tool() 955 | def analyze_image_from_path( 956 | image_path: str, 957 | model_name: str = "yolov8n.pt", 958 | confidence: float = 0.25, 959 | save_results: bool = False 960 | ) -> Dict[str, Any]: 961 | """ 962 | Return the YOLO CLI command for image analysis without executing it 963 | 964 | Args: 965 | image_path: Path to the image file 966 | model_name: YOLO model name 967 | confidence: Detection confidence threshold 968 | save_results: Whether to save results to disk 969 | 970 | Returns: 971 | Dictionary containing command that would be executed 972 | """ 973 | try: 974 | # Call detect_objects function with is_path=True 975 | return detect_objects( 976 | image_data=image_path, 977 | model_name=model_name, 978 | confidence=confidence, 979 | save_results=save_results, 980 | is_path=True 981 | ) 982 | except Exception as e: 983 | return { 984 | "error": f"Failed to generate analysis command: {str(e)}", 985 | "image_path": image_path 986 | } 987 | 988 | @mcp.tool() 989 | def test_connection() -> Dict[str, Any]: 990 | """ 991 | Test if YOLO CLI service is available 992 | 993 | Returns: 994 | Status information and available tools 995 | """ 996 | # Build a simple YOLO CLI version command 997 | cmd_args = ["--version"] 998 | result = run_yolo_cli(cmd_args) 999 | 1000 | return { 1001 | "status": "YOLO CLI command generator is running", 1002 | "command_mode": "Command generation only, no execution", 1003 | "version_command": result["command"], 1004 | "available_models": list_available_models(), 1005 | "available_tools": [ 1006 | "list_available_models", "detect_objects", "segment_objects", 1007 | "classify_image", "track_objects", "train_model", "validate_model", 1008 | "export_model", "start_camera_detection", "stop_camera_detection", 1009 | "get_camera_detections", "test_connection", 1010 | # Additional tools 1011 | "analyze_image_from_path", 1012 | "comprehensive_image_analysis" 1013 | ], 1014 | "note": "This service only generates YOLO commands without executing them" 1015 | } 1016 | 1017 | # Modify the main execution section 1018 | if __name__ == "__main__": 1019 | import platform 1020 | 1021 | logger.info("Starting YOLO CLI command generator service") 1022 | logger.info(f"Platform: {platform.system()} {platform.release()}") 1023 | logger.info("⚠️ Commands will be generated but NOT executed") 1024 | 1025 | # Initialize and run server 1026 | logger.info("Starting MCP server...") 1027 | mcp.run(transport='stdio') -------------------------------------------------------------------------------- /server.py: -------------------------------------------------------------------------------- 1 | # server.py 2 | import fnmatch 3 | import os 4 | import base64 5 | import cv2 6 | import time 7 | import threading 8 | from io import BytesIO 9 | from typing import List, Dict, Any, Optional, Union 10 | import numpy as np 11 | from PIL import Image 12 | 13 | from mcp.server.fastmcp import FastMCP 14 | from ultralytics import YOLO 15 | 16 | # Add this near the top of server.py with other imports 17 | import os.path 18 | import sys 19 | import logging 20 | import contextlib 21 | import logging 22 | import sys 23 | import contextlib 24 | import signal 25 | import atexit 26 | 27 | # Set up logging configuration - add this near the top of the file 28 | logging.basicConfig( 29 | level=logging.INFO, 30 | format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', 31 | handlers=[ 32 | logging.FileHandler("yolo_service.log"), 33 | logging.StreamHandler(sys.stderr) 34 | ] 35 | ) 36 | 37 | logger = logging.getLogger('yolo_service') 38 | 39 | # Global variables for camera control 40 | camera_running = False 41 | camera_thread = None 42 | detection_results = [] 43 | camera_last_access_time = 0 44 | CAMERA_INACTIVITY_TIMEOUT = 60 # Auto-shutdown after 60 seconds of inactivity 45 | 46 | @contextlib.contextmanager 47 | def redirect_stdout_to_stderr(): 48 | old_stdout = sys.stdout 49 | sys.stdout = sys.stderr 50 | try: 51 | yield 52 | finally: 53 | sys.stdout = old_stdout 54 | 55 | def camera_watchdog_thread(): 56 | """Monitor thread that auto-stops the camera after inactivity""" 57 | global camera_running, camera_last_access_time 58 | 59 | logger.info("Camera watchdog thread started") 60 | 61 | while True: 62 | # Sleep for a short time to avoid excessive CPU usage 63 | time.sleep(5) 64 | 65 | # Check if camera is running 66 | if camera_running: 67 | current_time = time.time() 68 | elapsed_time = current_time - camera_last_access_time 69 | 70 | # If no access for more than the timeout, auto-stop 71 | if elapsed_time > CAMERA_INACTIVITY_TIMEOUT: 72 | logger.info(f"Auto-stopping camera after {elapsed_time:.1f} seconds of inactivity") 73 | stop_camera_detection() 74 | else: 75 | # If camera is not running, no need to check frequently 76 | time.sleep(10) 77 | 78 | 79 | def load_image(image_source, is_path=False): 80 | """ 81 | Load image from file path or base64 data 82 | 83 | Args: 84 | image_source: File path or base64 encoded image data 85 | is_path: Whether image_source is a file path 86 | 87 | Returns: 88 | PIL Image object 89 | """ 90 | try: 91 | if is_path: 92 | # Load image from file path 93 | if os.path.exists(image_source): 94 | return Image.open(image_source) 95 | else: 96 | raise FileNotFoundError(f"Image file not found: {image_source}") 97 | else: 98 | # Load image from base64 data 99 | image_bytes = base64.b64decode(image_source) 100 | return Image.open(BytesIO(image_bytes)) 101 | except Exception as e: 102 | raise ValueError(f"Failed to load image: {str(e)}") 103 | 104 | # Create MCP server 105 | mcp = FastMCP("YOLO_Service") 106 | 107 | # Global model cache 108 | models = {} 109 | 110 | def get_model(model_name: str = "yolov8n.pt") -> YOLO: 111 | """Get or load YOLO model from any of the configured model directories""" 112 | if model_name in models: 113 | return models[model_name] 114 | 115 | # Try to find the model in any of the configured directories 116 | model_path = None 117 | for directory in CONFIG["model_dirs"]: 118 | potential_path = os.path.join(directory, model_name) 119 | if os.path.exists(potential_path): 120 | model_path = potential_path 121 | break 122 | 123 | if model_path is None: 124 | available = list_available_models() 125 | available_str = ", ".join(available) if available else "none" 126 | raise FileNotFoundError(f"Model '{model_name}' not found in any configured directories. Available models: {available_str}") 127 | 128 | # Load and cache the model - with stdout redirected 129 | logger.info(f"Loading model: {model_name} from {model_path}") 130 | with redirect_stdout_to_stderr(): 131 | models[model_name] = YOLO(model_path) 132 | return models[model_name] 133 | 134 | # Global configuration 135 | CONFIG = { 136 | "model_dirs": [ 137 | ".", # Current directory 138 | "./models", # Models subdirectory 139 | os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"), # Absolute path to models 140 | # Add any other potential model directories here 141 | ] 142 | } 143 | 144 | 145 | 146 | # Add a new tool to get information about model directories 147 | @mcp.tool() 148 | def get_model_directories() -> Dict[str, Any]: 149 | """Get information about configured model directories and available models""" 150 | directories = [] 151 | 152 | for directory in CONFIG["model_dirs"]: 153 | dir_info = { 154 | "path": directory, 155 | "exists": os.path.exists(directory), 156 | "is_directory": os.path.isdir(directory) if os.path.exists(directory) else False, 157 | "models": [] 158 | } 159 | 160 | if dir_info["exists"] and dir_info["is_directory"]: 161 | for filename in os.listdir(directory): 162 | if filename.endswith(".pt"): 163 | dir_info["models"].append(filename) 164 | 165 | directories.append(dir_info) 166 | 167 | return { 168 | "configured_directories": CONFIG["model_dirs"], 169 | "directory_details": directories, 170 | "available_models": list_available_models(), 171 | "loaded_models": list(models.keys()) 172 | } 173 | 174 | @mcp.tool() 175 | def detect_objects( 176 | image_data: str, 177 | model_name: str = "yolov8n.pt", 178 | confidence: float = 0.25, 179 | save_results: bool = False, 180 | is_path: bool = False 181 | ) -> Dict[str, Any]: 182 | """ 183 | Detect objects in an image using YOLO 184 | 185 | Args: 186 | image_data: Base64 encoded image or file path (if is_path=True) 187 | model_name: YOLO model name 188 | confidence: Detection confidence threshold 189 | save_results: Whether to save results to disk 190 | is_path: Whether image_data is a file path 191 | 192 | Returns: 193 | Dictionary containing detection results 194 | """ 195 | try: 196 | # Load image (supports path or base64) 197 | image = load_image(image_data, is_path=is_path) 198 | 199 | # Load model and perform detection - with stdout redirected 200 | model = get_model(model_name) 201 | with redirect_stdout_to_stderr(): # Ensure all YOLO outputs go to stderr 202 | results = model.predict(image, conf=confidence, save=save_results) 203 | 204 | # Format results 205 | formatted_results = [] 206 | for result in results: 207 | boxes = result.boxes 208 | detections = [] 209 | 210 | for i in range(len(boxes)): 211 | box = boxes[i] 212 | x1, y1, x2, y2 = box.xyxy[0].tolist() 213 | confidence = float(box.conf[0]) 214 | class_id = int(box.cls[0]) 215 | class_name = result.names[class_id] 216 | 217 | detections.append({ 218 | "box": [x1, y1, x2, y2], 219 | "confidence": confidence, 220 | "class_id": class_id, 221 | "class_name": class_name 222 | }) 223 | 224 | formatted_results.append({ 225 | "detections": detections, 226 | "image_shape": result.orig_shape 227 | }) 228 | 229 | return { 230 | "results": formatted_results, 231 | "model_used": model_name, 232 | "total_detections": sum(len(r["detections"]) for r in formatted_results), 233 | "source": image_data if is_path else "base64_image" 234 | } 235 | except Exception as e: 236 | logger.error(f"Error in detect_objects: {str(e)}") 237 | return { 238 | "error": f"Failed to detect objects: {str(e)}", 239 | "source": image_data if is_path else "base64_image" 240 | } 241 | 242 | @mcp.tool() 243 | def segment_objects( 244 | image_data: str, 245 | model_name: str = "yolov11n-seg.pt", 246 | confidence: float = 0.25, 247 | save_results: bool = False, 248 | is_path: bool = False 249 | ) -> Dict[str, Any]: 250 | """ 251 | Perform instance segmentation on an image using YOLO 252 | 253 | Args: 254 | image_data: Base64 encoded image or file path (if is_path=True) 255 | model_name: YOLO segmentation model name 256 | confidence: Detection confidence threshold 257 | save_results: Whether to save results to disk 258 | is_path: Whether image_data is a file path 259 | 260 | Returns: 261 | Dictionary containing segmentation results 262 | """ 263 | try: 264 | # Load image (supports path or base64) 265 | image = load_image(image_data, is_path=is_path) 266 | 267 | # Load model and perform segmentation 268 | model = get_model(model_name) 269 | with redirect_stdout_to_stderr(): # Add this context manager 270 | results = model.predict(image, conf=confidence, save=save_results) 271 | 272 | # Format results 273 | formatted_results = [] 274 | for result in results: 275 | if not hasattr(result, 'masks') or result.masks is None: 276 | continue 277 | 278 | boxes = result.boxes 279 | masks = result.masks 280 | segments = [] 281 | 282 | for i in range(len(boxes)): 283 | box = boxes[i] 284 | mask = masks[i].data[0].cpu().numpy() if masks else None 285 | 286 | x1, y1, x2, y2 = box.xyxy[0].tolist() 287 | confidence = float(box.conf[0]) 288 | class_id = int(box.cls[0]) 289 | class_name = result.names[class_id] 290 | 291 | segment = { 292 | "box": [x1, y1, x2, y2], 293 | "confidence": confidence, 294 | "class_id": class_id, 295 | "class_name": class_name 296 | } 297 | 298 | if mask is not None: 299 | # Convert binary mask to simplified format for API response 300 | segment["mask"] = mask.tolist() 301 | 302 | segments.append(segment) 303 | 304 | formatted_results.append({ 305 | "segments": segments, 306 | "image_shape": result.orig_shape 307 | }) 308 | 309 | return { 310 | "results": formatted_results, 311 | "model_used": model_name, 312 | "total_segments": sum(len(r["segments"]) for r in formatted_results), 313 | "source": image_data if is_path else "base64_image" 314 | } 315 | except Exception as e: 316 | return { 317 | "error": f"Failed to segment objects: {str(e)}", 318 | "source": image_data if is_path else "base64_image" 319 | } 320 | 321 | 322 | @mcp.tool() 323 | def classify_image( 324 | image_data: str, 325 | model_name: str = "yolov11n-cls.pt", 326 | top_k: int = 5, 327 | save_results: bool = False, 328 | is_path: bool = False 329 | ) -> Dict[str, Any]: 330 | """ 331 | Classify an image using YOLO classification model 332 | 333 | Args: 334 | image_data: Base64 encoded image or file path (if is_path=True) 335 | model_name: YOLO classification model name 336 | top_k: Number of top categories to return 337 | save_results: Whether to save results to disk 338 | is_path: Whether image_data is a file path 339 | 340 | Returns: 341 | Dictionary containing classification results 342 | """ 343 | try: 344 | # Load image (supports path or base64) 345 | image = load_image(image_data, is_path=is_path) 346 | 347 | # Load model and perform classification 348 | model = get_model(model_name) 349 | with redirect_stdout_to_stderr(): # Add this context manager 350 | results = model.predict(image, save=save_results) 351 | 352 | # Format results 353 | formatted_results = [] 354 | for result in results: 355 | if not hasattr(result, 'probs') or result.probs is None: 356 | continue 357 | 358 | probs = result.probs 359 | top_indices = probs.top5 360 | top_probs = probs.top5conf.tolist() 361 | top_classes = [result.names[idx] for idx in top_indices] 362 | 363 | classifications = [ 364 | {"class_id": int(idx), "class_name": name, "probability": float(prob)} 365 | for idx, name, prob in zip(top_indices[:top_k], top_classes[:top_k], top_probs[:top_k]) 366 | ] 367 | 368 | formatted_results.append({ 369 | "classifications": classifications, 370 | "image_shape": result.orig_shape 371 | }) 372 | 373 | return { 374 | "results": formatted_results, 375 | "model_used": model_name, 376 | "top_k": top_k, 377 | "source": image_data if is_path else "base64_image" 378 | } 379 | except Exception as e: 380 | return { 381 | "error": f"Failed to classify image: {str(e)}", 382 | "source": image_data if is_path else "base64_image" 383 | } 384 | 385 | 386 | @mcp.tool() 387 | def track_objects( 388 | image_data: str, 389 | model_name: str = "yolov8n.pt", 390 | confidence: float = 0.25, 391 | tracker: str = "bytetrack.yaml", 392 | save_results: bool = False 393 | ) -> Dict[str, Any]: 394 | """ 395 | Track objects in an image sequence using YOLO 396 | 397 | Args: 398 | image_data: Base64 encoded image 399 | model_name: YOLO model name 400 | confidence: Detection confidence threshold 401 | tracker: Tracker name to use (e.g., 'bytetrack.yaml', 'botsort.yaml') 402 | save_results: Whether to save results to disk 403 | 404 | Returns: 405 | Dictionary containing tracking results 406 | """ 407 | # Decode Base64 image 408 | image_bytes = base64.b64decode(image_data) 409 | image = Image.open(BytesIO(image_bytes)) 410 | 411 | # Load model and perform tracking 412 | model = get_model(model_name) 413 | # Add redirect_stdout_to_stderr context manager 414 | with redirect_stdout_to_stderr(): 415 | results = model.track(image, conf=confidence, tracker=tracker, save=save_results) 416 | 417 | # Format results 418 | formatted_results = [] 419 | for result in results: 420 | if not hasattr(result, 'boxes') or result.boxes is None: 421 | continue 422 | 423 | boxes = result.boxes 424 | tracks = [] 425 | 426 | for i in range(len(boxes)): 427 | box = boxes[i] 428 | x1, y1, x2, y2 = box.xyxy[0].tolist() 429 | confidence = float(box.conf[0]) 430 | class_id = int(box.cls[0]) 431 | class_name = result.names[class_id] 432 | 433 | # Extract track ID (if any) 434 | track_id = int(box.id[0]) if box.id is not None else None 435 | 436 | track = { 437 | "box": [x1, y1, x2, y2], 438 | "confidence": confidence, 439 | "class_id": class_id, 440 | "class_name": class_name, 441 | "track_id": track_id 442 | } 443 | 444 | tracks.append(track) 445 | 446 | formatted_results.append({ 447 | "tracks": tracks, 448 | "image_shape": result.orig_shape 449 | }) 450 | 451 | return { 452 | "results": formatted_results, 453 | "model_used": model_name, 454 | "tracker": tracker, 455 | "total_tracks": sum(len(r["tracks"]) for r in formatted_results) 456 | } 457 | 458 | # 3. FIX train_model FUNCTION TO USE REDIRECTION: 459 | @mcp.tool() 460 | def train_model( 461 | dataset_path: str, 462 | model_name: str = "yolov8n.pt", 463 | epochs: int = 100, 464 | imgsz: int = 640, 465 | batch: int = 16, 466 | name: str = "yolo_custom_model", 467 | project: str = "runs/train" 468 | ) -> Dict[str, Any]: 469 | """ 470 | Train a YOLO model on a custom dataset 471 | 472 | Args: 473 | dataset_path: Path to YOLO format dataset 474 | model_name: Base model to start with 475 | epochs: Number of training epochs 476 | imgsz: Image size for training 477 | batch: Batch size 478 | name: Name for the training run 479 | project: Project directory 480 | 481 | Returns: 482 | Dictionary containing training results 483 | """ 484 | # Validate dataset path 485 | if not os.path.exists(dataset_path): 486 | return {"error": f"Dataset not found: {dataset_path}"} 487 | 488 | # Initialize model 489 | model = get_model(model_name) 490 | 491 | # Train model 492 | try: 493 | # Add redirect_stdout_to_stderr context manager 494 | with redirect_stdout_to_stderr(): 495 | results = model.train( 496 | data=dataset_path, 497 | epochs=epochs, 498 | imgsz=imgsz, 499 | batch=batch, 500 | name=name, 501 | project=project 502 | ) 503 | 504 | # Get best model path 505 | best_model_path = os.path.join(project, name, "weights", "best.pt") 506 | 507 | return { 508 | "status": "success", 509 | "model_path": best_model_path, 510 | "epochs_completed": epochs, 511 | "final_metrics": { 512 | "precision": float(results.results_dict.get("metrics/precision(B)", 0)), 513 | "recall": float(results.results_dict.get("metrics/recall(B)", 0)), 514 | "mAP50": float(results.results_dict.get("metrics/mAP50(B)", 0)), 515 | "mAP50-95": float(results.results_dict.get("metrics/mAP50-95(B)", 0)) 516 | } 517 | } 518 | except Exception as e: 519 | return {"error": f"Training failed: {str(e)}"} 520 | 521 | # 4. FIX validate_model FUNCTION TO USE REDIRECTION: 522 | @mcp.tool() 523 | def validate_model( 524 | model_path: str, 525 | data_path: str, 526 | imgsz: int = 640, 527 | batch: int = 16 528 | ) -> Dict[str, Any]: 529 | """ 530 | Validate a YOLO model on a dataset 531 | 532 | Args: 533 | model_path: Path to YOLO model (.pt file) 534 | data_path: Path to YOLO format validation dataset 535 | imgsz: Image size for validation 536 | batch: Batch size 537 | 538 | Returns: 539 | Dictionary containing validation results 540 | """ 541 | # Validate model path 542 | if not os.path.exists(model_path): 543 | return {"error": f"Model file not found: {model_path}"} 544 | 545 | # Validate dataset path 546 | if not os.path.exists(data_path): 547 | return {"error": f"Dataset not found: {data_path}"} 548 | 549 | # Load model 550 | try: 551 | model = get_model(model_path) 552 | except Exception as e: 553 | return {"error": f"Failed to load model: {str(e)}"} 554 | 555 | # Validate model 556 | try: 557 | # Add redirect_stdout_to_stderr context manager 558 | with redirect_stdout_to_stderr(): 559 | results = model.val(data=data_path, imgsz=imgsz, batch=batch) 560 | 561 | return { 562 | "status": "success", 563 | "metrics": { 564 | "precision": float(results.results_dict.get("metrics/precision(B)", 0)), 565 | "recall": float(results.results_dict.get("metrics/recall(B)", 0)), 566 | "mAP50": float(results.results_dict.get("metrics/mAP50(B)", 0)), 567 | "mAP50-95": float(results.results_dict.get("metrics/mAP50-95(B)", 0)) 568 | } 569 | } 570 | except Exception as e: 571 | return {"error": f"Validation failed: {str(e)}"} 572 | 573 | # 5. FIX export_model FUNCTION TO USE REDIRECTION: 574 | @mcp.tool() 575 | def export_model( 576 | model_path: str, 577 | format: str = "onnx", 578 | imgsz: int = 640 579 | ) -> Dict[str, Any]: 580 | """ 581 | Export a YOLO model to different formats 582 | 583 | Args: 584 | model_path: Path to YOLO model (.pt file) 585 | format: Export format (onnx, torchscript, openvino, etc.) 586 | imgsz: Image size for export 587 | 588 | Returns: 589 | Dictionary containing export results 590 | """ 591 | # Validate model path 592 | if not os.path.exists(model_path): 593 | return {"error": f"Model file not found: {model_path}"} 594 | 595 | # Valid export formats 596 | valid_formats = [ 597 | "torchscript", "onnx", "openvino", "engine", "coreml", "saved_model", 598 | "pb", "tflite", "edgetpu", "tfjs", "paddle" 599 | ] 600 | 601 | if format not in valid_formats: 602 | return {"error": f"Invalid export format: {format}. Valid formats include: {', '.join(valid_formats)}"} 603 | 604 | # Load model 605 | try: 606 | model = get_model(model_path) 607 | except Exception as e: 608 | return {"error": f"Failed to load model: {str(e)}"} 609 | 610 | # Export model 611 | try: 612 | # Add redirect_stdout_to_stderr context manager 613 | with redirect_stdout_to_stderr(): 614 | export_path = model.export(format=format, imgsz=imgsz) 615 | 616 | return { 617 | "status": "success", 618 | "export_path": str(export_path), 619 | "format": format 620 | } 621 | except Exception as e: 622 | return {"error": f"Export failed: {str(e)}"} 623 | 624 | # 6. ADD REDIRECTION TO get_model_info FUNCTION: 625 | @mcp.resource("model_info/{model_name}") 626 | def get_model_info(model_name: str) -> Dict[str, Any]: 627 | """ 628 | Get information about a YOLO model 629 | 630 | Args: 631 | model_name: YOLO model name 632 | 633 | Returns: 634 | Dictionary containing model information 635 | """ 636 | try: 637 | model = get_model(model_name) 638 | 639 | # Get model task 640 | task = 'detect' # Default task 641 | if 'seg' in model_name: 642 | task = 'segment' 643 | elif 'pose' in model_name: 644 | task = 'pose' 645 | elif 'cls' in model_name: 646 | task = 'classify' 647 | elif 'obb' in model_name: 648 | task = 'obb' 649 | 650 | # Make sure any model property access that might trigger output is wrapped 651 | with redirect_stdout_to_stderr(): 652 | yaml_str = str(model.yaml) 653 | pt_path = str(model.pt_path) if hasattr(model, 'pt_path') else None 654 | class_names = model.names 655 | 656 | # Get model info 657 | return { 658 | "model_name": model_name, 659 | "task": task, 660 | "yaml": yaml_str, 661 | "pt_path": pt_path, 662 | "class_names": class_names 663 | } 664 | except Exception as e: 665 | return {"error": f"Failed to get model info: {str(e)}"} 666 | 667 | # 7. MODIFY list_available_models to use logging instead of print 668 | @mcp.tool() 669 | def list_available_models() -> List[str]: 670 | """List available YOLO models that actually exist on disk in any configured directory""" 671 | # Common YOLO model patterns 672 | model_patterns = [ 673 | "yolov11*.pt", 674 | "yolov8*.pt" 675 | ] 676 | 677 | # Find all existing models in all configured directories 678 | available_models = set() 679 | for directory in CONFIG["model_dirs"]: 680 | if not os.path.exists(directory): 681 | continue 682 | 683 | # Check for model files directly 684 | for filename in os.listdir(directory): 685 | if filename.endswith(".pt") and any( 686 | fnmatch.fnmatch(filename, pattern) for pattern in model_patterns 687 | ): 688 | available_models.add(filename) 689 | 690 | # Convert to sorted list 691 | result = sorted(list(available_models)) 692 | 693 | if not result: 694 | # Replace print with logger 695 | logger.warning("No model files found in configured directories.") 696 | return ["No models available - download models to any of these directories: " + ", ".join(CONFIG["model_dirs"])] 697 | 698 | return result 699 | @mcp.resource("model_info/{model_name}") 700 | def get_model_info(model_name: str) -> Dict[str, Any]: 701 | """ 702 | Get information about a YOLO model 703 | 704 | Args: 705 | model_name: YOLO model name 706 | 707 | Returns: 708 | Dictionary containing model information 709 | """ 710 | try: 711 | model = get_model(model_name) 712 | 713 | # Get model task 714 | task = 'detect' # Default task 715 | if 'seg' in model_name: 716 | task = 'segment' 717 | elif 'pose' in model_name: 718 | task = 'pose' 719 | elif 'cls' in model_name: 720 | task = 'classify' 721 | elif 'obb' in model_name: 722 | task = 'obb' 723 | 724 | # Get model info 725 | return { 726 | "model_name": model_name, 727 | "task": task, 728 | "yaml": str(model.yaml), 729 | "pt_path": str(model.pt_path) if hasattr(model, 'pt_path') else None, 730 | "class_names": model.names 731 | } 732 | except Exception as e: 733 | return {"error": f"Failed to get model info: {str(e)}"} 734 | 735 | @mcp.tool() 736 | def list_available_models() -> List[str]: 737 | """List available YOLO models that actually exist on disk in any configured directory""" 738 | # Common YOLO model patterns 739 | model_patterns = [ 740 | "yolov11*.pt", 741 | "yolov8*.pt" 742 | ] 743 | 744 | # Find all existing models in all configured directories 745 | available_models = set() 746 | for directory in CONFIG["model_dirs"]: 747 | if not os.path.exists(directory): 748 | continue 749 | 750 | # Check for model files directly 751 | for filename in os.listdir(directory): 752 | if filename.endswith(".pt") and any( 753 | fnmatch.fnmatch(filename, pattern) for pattern in model_patterns 754 | ): 755 | available_models.add(filename) 756 | 757 | # Convert to sorted list 758 | result = sorted(list(available_models)) 759 | 760 | if not result: 761 | print("Warning: No model files found in configured directories.") 762 | return ["No models available - download models to any of these directories: " + ", ".join(CONFIG["model_dirs"])] 763 | 764 | return result 765 | 766 | 767 | 768 | # Camera detection background thread 769 | camera_thread = None 770 | camera_running = False 771 | detection_results = [] 772 | 773 | def camera_detection_thread(model_name, confidence, fps_limit=30, camera_id=0): 774 | """Background thread for camera detection""" 775 | global camera_running, detection_results 776 | 777 | # Load model 778 | try: 779 | with redirect_stdout_to_stderr(): 780 | model = get_model(model_name) 781 | logger.info(f"Model {model_name} loaded successfully") 782 | except Exception as e: 783 | logger.error(f"Error loading model: {str(e)}") 784 | camera_running = False 785 | detection_results.append({ 786 | "timestamp": time.time(), 787 | "error": f"Failed to load model: {str(e)}", 788 | "detections": [] 789 | }) 790 | return 791 | 792 | # Rest of the function... 793 | # Try to open camera with multiple attempts and multiple camera IDs if necessary 794 | cap = None 795 | error_message = "" 796 | 797 | # Try camera IDs from 0 to 2 798 | for cam_id in range(3): 799 | try: 800 | logger.info(f"Attempting to open camera with ID {cam_id}...") 801 | cap = cv2.VideoCapture(cam_id) 802 | if cap.isOpened(): 803 | logger.info(f"Successfully opened camera {cam_id}") 804 | break 805 | except Exception as e: 806 | error_message = f"Error opening camera {cam_id}: {str(e)}" 807 | logger.error(error_message) 808 | 809 | # Check if any camera was successfully opened 810 | if cap is None or not cap.isOpened(): 811 | logger.error("Error: Could not open any camera.") 812 | camera_running = False 813 | detection_results.append({ 814 | "timestamp": time.time(), 815 | "error": "Failed to open camera. Make sure camera is connected and not in use by another application.", 816 | "camera_status": "unavailable", 817 | "detections": [] 818 | }) 819 | return 820 | 821 | # Get camera properties for diagnostics 822 | width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) 823 | height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) 824 | fps = cap.get(cv2.CAP_PROP_FPS) 825 | 826 | logger.info(f"Camera properties: {width}x{height} at {fps} FPS") 827 | 828 | # Calculate frame interval based on fps_limit 829 | frame_interval = 1.0 / fps_limit 830 | frame_count = 0 831 | error_count = 0 832 | 833 | while camera_running: 834 | start_time = time.time() 835 | 836 | try: 837 | ret, frame = cap.read() 838 | if not ret: 839 | logger.warning(f"Error: Failed to capture frame (attempt {error_count+1}).") 840 | error_count += 1 841 | 842 | # Add error to detection results 843 | detection_results.append({ 844 | "timestamp": time.time(), 845 | "error": f"Failed to capture frame (attempt {error_count})", 846 | "camera_status": "error", 847 | "detections": [] 848 | }) 849 | 850 | # If we have consistent failures, try to restart the camera 851 | if error_count >= 5: 852 | logger.warning("Too many frame capture errors, attempting to restart camera...") 853 | cap.release() 854 | time.sleep(1) 855 | cap = cv2.VideoCapture(camera_id) 856 | error_count = 0 857 | if not cap.isOpened(): 858 | logger.error("Failed to reopen camera after errors.") 859 | break 860 | 861 | time.sleep(1) # Wait before trying again 862 | continue 863 | 864 | # Reset error count on successful frame capture 865 | error_count = 0 866 | frame_count += 1 867 | 868 | # Perform detection on frame 869 | with redirect_stdout_to_stderr(): # Add this context manager 870 | results = model.predict(frame, conf=confidence) 871 | 872 | # Update detection results (only keep the last 10) 873 | if len(detection_results) >= 10: 874 | detection_results.pop(0) 875 | 876 | # Format results 877 | for result in results: 878 | boxes = result.boxes 879 | detections = [] 880 | 881 | for i in range(len(boxes)): 882 | box = boxes[i] 883 | x1, y1, x2, y2 = box.xyxy[0].tolist() 884 | conf = float(box.conf[0]) 885 | class_id = int(box.cls[0]) 886 | class_name = result.names[class_id] 887 | 888 | detections.append({ 889 | "box": [x1, y1, x2, y2], 890 | "confidence": conf, 891 | "class_id": class_id, 892 | "class_name": class_name 893 | }) 894 | 895 | detection_results.append({ 896 | "timestamp": time.time(), 897 | "frame_count": frame_count, 898 | "detections": detections, 899 | "camera_status": "running", 900 | "image_shape": result.orig_shape 901 | }) 902 | 903 | # Log occasional status 904 | if frame_count % 30 == 0: 905 | logger.info(f"Camera running: processed {frame_count} frames") 906 | detection_count = sum(len(r.get("detections", [])) for r in detection_results if "detections" in r) 907 | logger.info(f"Total detections in current buffer: {detection_count}") 908 | 909 | # Limit FPS by waiting if necessary 910 | elapsed = time.time() - start_time 911 | if elapsed < frame_interval: 912 | time.sleep(frame_interval - elapsed) 913 | 914 | except Exception as e: 915 | logger.error(f"Error in camera thread: {str(e)}") 916 | detection_results.append({ 917 | "timestamp": time.time(), 918 | "error": f"Exception in camera processing: {str(e)}", 919 | "camera_status": "error", 920 | "detections": [] 921 | }) 922 | time.sleep(1) # Wait before continuing 923 | 924 | # Clean up 925 | logger.info("Shutting down camera...") 926 | if cap is not None: 927 | cap.release() 928 | 929 | @mcp.tool() 930 | def start_camera_detection( 931 | model_name: str = "yolov8n.pt", 932 | confidence: float = 0.25, 933 | camera_id: int = 0 934 | ) -> Dict[str, Any]: 935 | """ 936 | Start realtime object detection using the computer's camera 937 | 938 | Args: 939 | model_name: YOLO model name to use 940 | confidence: Detection confidence threshold 941 | camera_id: Camera device ID (0 is usually the default camera) 942 | 943 | Returns: 944 | Status of camera detection 945 | """ 946 | global camera_thread, camera_running, detection_results, camera_last_access_time 947 | 948 | # Check if already running 949 | if camera_running: 950 | # Update last access time 951 | camera_last_access_time = time.time() 952 | return {"status": "success", "message": "Camera detection is already running"} 953 | 954 | # Clear previous results 955 | detection_results = [] 956 | 957 | # First, try to check if OpenCV is properly installed 958 | try: 959 | cv2_version = cv2.__version__ 960 | logger.info(f"OpenCV version: {cv2_version}") 961 | except Exception as e: 962 | logger.error(f"OpenCV not properly installed: {str(e)}") 963 | return { 964 | "status": "error", 965 | "message": f"OpenCV not properly installed: {str(e)}", 966 | "solution": "Please check OpenCV installation" 967 | } 968 | 969 | # Start detection thread 970 | camera_running = True 971 | camera_last_access_time = time.time() # Update access time 972 | camera_thread = threading.Thread( 973 | target=camera_detection_thread, 974 | args=(model_name, confidence, 30, camera_id), 975 | daemon=True 976 | ) 977 | camera_thread.start() 978 | 979 | # Add initial status to detection results 980 | detection_results.append({ 981 | "timestamp": time.time(), 982 | "system_info": { 983 | "os": platform.system() if 'platform' in globals() else "Unknown", 984 | "opencv_version": cv2.__version__, 985 | "camera_id": camera_id 986 | }, 987 | "camera_status": "starting", 988 | "detections": [] 989 | }) 990 | 991 | return { 992 | "status": "success", 993 | "message": f"Started camera detection using model {model_name}", 994 | "model": model_name, 995 | "confidence": confidence, 996 | "camera_id": camera_id, 997 | "auto_shutdown": f"Camera will auto-shutdown after {CAMERA_INACTIVITY_TIMEOUT} seconds of inactivity", 998 | "note": "If camera doesn't work, try different camera_id values (0, 1, or 2)" 999 | } 1000 | 1001 | 1002 | @mcp.tool() 1003 | def stop_camera_detection() -> Dict[str, Any]: 1004 | """ 1005 | Stop realtime camera detection 1006 | 1007 | Returns: 1008 | Status message 1009 | """ 1010 | global camera_running 1011 | 1012 | if not camera_running: 1013 | return {"status": "error", "message": "Camera detection is not running"} 1014 | 1015 | logger.info("Stopping camera detection by user request") 1016 | camera_running = False 1017 | 1018 | # Wait for thread to terminate 1019 | if camera_thread and camera_thread.is_alive(): 1020 | camera_thread.join(timeout=2.0) 1021 | 1022 | return { 1023 | "status": "success", 1024 | "message": "Stopped camera detection" 1025 | } 1026 | 1027 | 1028 | @mcp.tool() 1029 | def get_camera_detections() -> Dict[str, Any]: 1030 | """ 1031 | Get the latest detections from the camera 1032 | 1033 | Returns: 1034 | Dictionary with recent detections 1035 | """ 1036 | global detection_results, camera_thread, camera_last_access_time 1037 | 1038 | # Update the last access time whenever this function is called 1039 | if camera_running: 1040 | camera_last_access_time = time.time() 1041 | 1042 | # Check if thread is alive 1043 | thread_alive = camera_thread is not None and camera_thread.is_alive() 1044 | 1045 | # If camera_running is True but thread is dead, there's an issue 1046 | if camera_running and not thread_alive: 1047 | return { 1048 | "status": "error", 1049 | "message": "Camera thread has stopped unexpectedly", 1050 | "is_running": False, 1051 | "camera_status": "error", 1052 | "thread_alive": thread_alive, 1053 | "detections": detection_results, 1054 | "count": len(detection_results), 1055 | "solution": "Please try restart the camera with a different camera_id" 1056 | } 1057 | 1058 | if not camera_running: 1059 | return { 1060 | "status": "error", 1061 | "message": "Camera detection is not running", 1062 | "is_running": False, 1063 | "camera_status": "stopped" 1064 | } 1065 | 1066 | # Check for errors in detection results 1067 | errors = [result.get("error") for result in detection_results if "error" in result] 1068 | recent_errors = errors[-5:] if errors else [] 1069 | 1070 | # Count actual detections 1071 | detection_count = sum(len(result.get("detections", [])) for result in detection_results if "detections" in result) 1072 | 1073 | return { 1074 | "status": "success", 1075 | "is_running": camera_running, 1076 | "thread_alive": thread_alive, 1077 | "detections": detection_results, 1078 | "count": len(detection_results), 1079 | "total_detections": detection_count, 1080 | "recent_errors": recent_errors if recent_errors else None, 1081 | "camera_status": "error" if recent_errors else "running", 1082 | "inactivity_timeout": { 1083 | "seconds_remaining": int(CAMERA_INACTIVITY_TIMEOUT - (time.time() - camera_last_access_time)), 1084 | "last_access": camera_last_access_time 1085 | } 1086 | } 1087 | 1088 | def cleanup_resources(): 1089 | """Clean up resources when the server is shutting down""" 1090 | global camera_running 1091 | 1092 | logger.info("Cleaning up resources...") 1093 | 1094 | # Stop camera if it's running 1095 | if camera_running: 1096 | logger.info("Shutting down camera during server exit") 1097 | camera_running = False 1098 | 1099 | # Give the camera thread a moment to clean up 1100 | if camera_thread and camera_thread.is_alive(): 1101 | camera_thread.join(timeout=2.0) 1102 | 1103 | logger.info("Cleanup complete") 1104 | atexit.register(cleanup_resources) 1105 | 1106 | def signal_handler(sig, frame): 1107 | """Handle termination signals""" 1108 | logger.info(f"Received signal {sig}, shutting down...") 1109 | cleanup_resources() 1110 | sys.exit(0) 1111 | 1112 | signal.signal(signal.SIGINT, signal_handler) 1113 | signal.signal(signal.SIGTERM, signal_handler) 1114 | 1115 | def start_watchdog(): 1116 | """Start the camera watchdog thread""" 1117 | watchdog = threading.Thread( 1118 | target=camera_watchdog_thread, 1119 | daemon=True 1120 | ) 1121 | watchdog.start() 1122 | return watchdog 1123 | 1124 | @mcp.tool() 1125 | def comprehensive_image_analysis( 1126 | image_path: str, 1127 | confidence: float = 0.25, 1128 | save_results: bool = False 1129 | ) -> Dict[str, Any]: 1130 | """ 1131 | Perform comprehensive analysis on an image by combining multiple model results 1132 | 1133 | Args: 1134 | image_path: Path to the image file 1135 | confidence: Detection confidence threshold 1136 | save_results: Whether to save results to disk 1137 | 1138 | Returns: 1139 | Dictionary containing comprehensive analysis results 1140 | """ 1141 | try: 1142 | if not os.path.exists(image_path): 1143 | return {"error": f"Image file not found: {image_path}"} 1144 | 1145 | # Load image 1146 | image = load_image(image_path, is_path=True) 1147 | 1148 | analysis_results = {} 1149 | 1150 | # 1. Object detection 1151 | object_model = get_model("yolov11n.pt") 1152 | with redirect_stdout_to_stderr(): # Add this context manager 1153 | object_results = object_model.predict(image, conf=confidence, save=save_results) 1154 | 1155 | # Process object detection results 1156 | detected_objects = [] 1157 | for result in object_results: 1158 | boxes = result.boxes 1159 | for i in range(len(boxes)): 1160 | box = boxes[i] 1161 | conf = float(box.conf[0]) 1162 | class_id = int(box.cls[0]) 1163 | class_name = result.names[class_id] 1164 | detected_objects.append({ 1165 | "class_name": class_name, 1166 | "confidence": conf 1167 | }) 1168 | analysis_results["objects"] = detected_objects 1169 | 1170 | # 2. Scene classification 1171 | try: 1172 | cls_model = get_model("yolov8n-cls.pt") 1173 | with redirect_stdout_to_stderr(): # Add this context manager 1174 | cls_results = cls_model.predict(image, save=False) 1175 | 1176 | scene_classifications = [] 1177 | for result in cls_results: 1178 | if hasattr(result, 'probs') and result.probs is not None: 1179 | probs = result.probs 1180 | top_indices = probs.top5 1181 | top_probs = probs.top5conf.tolist() 1182 | top_classes = [result.names[idx] for idx in top_indices] 1183 | 1184 | for idx, name, prob in zip(top_indices[:3], top_classes[:3], top_probs[:3]): 1185 | scene_classifications.append({ 1186 | "class_name": name, 1187 | "probability": float(prob) 1188 | }) 1189 | analysis_results["scene"] = scene_classifications 1190 | except Exception as e: 1191 | analysis_results["scene_error"] = str(e) 1192 | 1193 | # 3. Human pose detection 1194 | try: 1195 | pose_model = get_model("yolov8n-pose.pt") 1196 | with redirect_stdout_to_stderr(): # Add this context manager 1197 | pose_results = pose_model.predict(image, conf=confidence, save=False) 1198 | 1199 | detected_poses = [] 1200 | for result in pose_results: 1201 | if hasattr(result, 'keypoints') and result.keypoints is not None: 1202 | boxes = result.boxes 1203 | keypoints = result.keypoints 1204 | 1205 | for i in range(len(boxes)): 1206 | box = boxes[i] 1207 | conf = float(box.conf[0]) 1208 | detected_poses.append({ 1209 | "person_confidence": conf, 1210 | "has_keypoints": keypoints[i].data.shape[1] if keypoints else 0 1211 | }) 1212 | analysis_results["poses"] = detected_poses 1213 | except Exception as e: 1214 | analysis_results["pose_error"] = str(e) 1215 | 1216 | # Rest of the function remains the same... 1217 | # 4. Comprehensive task description 1218 | tasks = [] 1219 | 1220 | # Detect main objects 1221 | main_objects = [obj["class_name"] for obj in detected_objects if obj["confidence"] > 0.5] 1222 | if "person" in main_objects: 1223 | tasks.append("Person Detection") 1224 | 1225 | # Check for weapon objects 1226 | weapon_objects = ["sword", "knife", "katana", "gun", "pistol", "rifle"] 1227 | weapons = [obj for obj in main_objects if any(weapon in obj.lower() for weapon in weapon_objects)] 1228 | if weapons: 1229 | tasks.append(f"Weapon Detection ({', '.join(weapons)})") 1230 | 1231 | # Count people 1232 | person_count = main_objects.count("person") 1233 | if person_count > 0: 1234 | tasks.append(f"Person Count ({person_count} people)") 1235 | 1236 | # Pose analysis 1237 | if "poses" in analysis_results and analysis_results["poses"]: 1238 | tasks.append("Human Pose Analysis") 1239 | 1240 | # Scene classification 1241 | if "scene" in analysis_results and analysis_results["scene"]: 1242 | scene_types = [scene["class_name"] for scene in analysis_results["scene"][:2]] 1243 | tasks.append(f"Scene Classification ({', '.join(scene_types)})") 1244 | 1245 | analysis_results["identified_tasks"] = tasks 1246 | 1247 | # Return comprehensive results 1248 | return { 1249 | "status": "success", 1250 | "image_path": image_path, 1251 | "analysis": analysis_results, 1252 | "summary": "Tasks identified in the image: " + ", ".join(tasks) if tasks else "No clear tasks identified" 1253 | } 1254 | except Exception as e: 1255 | return { 1256 | "status": "error", 1257 | "image_path": image_path, 1258 | "error": f"Comprehensive analysis failed: {str(e)}" 1259 | } 1260 | 1261 | 1262 | @mcp.tool() 1263 | def analyze_image_from_path( 1264 | image_path: str, 1265 | model_name: str = "yolov8n.pt", 1266 | confidence: float = 0.25, 1267 | save_results: bool = False 1268 | ) -> Dict[str, Any]: 1269 | """ 1270 | Analyze image from file path using YOLO 1271 | 1272 | Args: 1273 | image_path: Path to the image file 1274 | model_name: YOLO model name 1275 | confidence: Detection confidence threshold 1276 | save_results: Whether to save results to disk 1277 | 1278 | Returns: 1279 | Dictionary containing detection results 1280 | """ 1281 | try: 1282 | # Call detect_objects function with is_path=True 1283 | return detect_objects( 1284 | image_data=image_path, 1285 | model_name=model_name, 1286 | confidence=confidence, 1287 | save_results=save_results, 1288 | is_path=True 1289 | ) 1290 | except Exception as e: 1291 | return { 1292 | "error": f"Failed to analyze image: {str(e)}", 1293 | "image_path": image_path 1294 | } 1295 | 1296 | @mcp.tool() 1297 | def test_connection() -> Dict[str, Any]: 1298 | """ 1299 | Test if YOLO MCP service is running properly 1300 | 1301 | Returns: 1302 | Status information and available tools 1303 | """ 1304 | return { 1305 | "status": "YOLO MCP service is running normally", 1306 | "available_models": list_available_models(), 1307 | "available_tools": [ 1308 | "list_available_models", "detect_objects", "segment_objects", 1309 | "classify_image", "detect_poses", "detect_oriented_objects", 1310 | "track_objects", "train_model", "validate_model", 1311 | "export_model", "start_camera_detection", "stop_camera_detection", 1312 | "get_camera_detections", "test_connection", 1313 | # Additional tools 1314 | "analyze_image_from_path", 1315 | "comprehensive_image_analysis" 1316 | ], 1317 | "new_features": [ 1318 | "Support for loading images directly from file paths", 1319 | "Support for comprehensive image analysis with task identification", 1320 | "All detection functions support both file paths and base64 data" 1321 | ] 1322 | } 1323 | 1324 | # Modify the main execution section 1325 | if __name__ == "__main__": 1326 | logger.info("Starting YOLO MCP service") 1327 | 1328 | # Start the camera watchdog thread 1329 | watchdog_thread = start_watchdog() 1330 | 1331 | # Initialize and run server 1332 | mcp.run(transport='stdio') -------------------------------------------------------------------------------- /server_cli.py: -------------------------------------------------------------------------------- 1 | # server.py - CLI version 2 | import fnmatch 3 | import os 4 | import base64 5 | import cv2 6 | import time 7 | import threading 8 | import subprocess 9 | import json 10 | import tempfile 11 | import platform 12 | from io import BytesIO 13 | from typing import List, Dict, Any, Optional, Union 14 | import numpy as np 15 | from PIL import Image 16 | 17 | from mcp.server.fastmcp import FastMCP 18 | 19 | # Set up logging configuration 20 | import os.path 21 | import sys 22 | import logging 23 | import contextlib 24 | import signal 25 | import atexit 26 | 27 | logging.basicConfig( 28 | level=logging.INFO, 29 | format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', 30 | handlers=[ 31 | logging.FileHandler("yolo_service.log"), 32 | logging.StreamHandler(sys.stderr) 33 | ] 34 | ) 35 | camera_startup_status = None # Will store error details if startup fails 36 | camera_last_error = None 37 | logger = logging.getLogger('yolo_service') 38 | 39 | # Global variables for camera control 40 | camera_running = False 41 | camera_thread = None 42 | detection_results = [] 43 | camera_last_access_time = 0 44 | CAMERA_INACTIVITY_TIMEOUT = 60 # Auto-shutdown after 60 seconds of inactivity 45 | 46 | def camera_watchdog_thread(): 47 | """Monitor thread that auto-stops the camera after inactivity""" 48 | global camera_running, camera_last_access_time 49 | 50 | logger.info("Camera watchdog thread started") 51 | 52 | while True: 53 | # Sleep for a short time to avoid excessive CPU usage 54 | time.sleep(5) 55 | 56 | # Check if camera is running 57 | if camera_running: 58 | current_time = time.time() 59 | elapsed_time = current_time - camera_last_access_time 60 | 61 | # If no access for more than the timeout, auto-stop 62 | if elapsed_time > CAMERA_INACTIVITY_TIMEOUT: 63 | logger.info(f"Auto-stopping camera after {elapsed_time:.1f} seconds of inactivity") 64 | stop_camera_detection() 65 | else: 66 | # If camera is not running, no need to check frequently 67 | time.sleep(10) 68 | 69 | def load_image(image_source, is_path=False): 70 | """ 71 | Load image from file path or base64 data 72 | 73 | Args: 74 | image_source: File path or base64 encoded image data 75 | is_path: Whether image_source is a file path 76 | 77 | Returns: 78 | PIL Image object 79 | """ 80 | try: 81 | if is_path: 82 | # Load image from file path 83 | if os.path.exists(image_source): 84 | return Image.open(image_source) 85 | else: 86 | raise FileNotFoundError(f"Image file not found: {image_source}") 87 | else: 88 | # Load image from base64 data 89 | image_bytes = base64.b64decode(image_source) 90 | return Image.open(BytesIO(image_bytes)) 91 | except Exception as e: 92 | raise ValueError(f"Failed to load image: {str(e)}") 93 | 94 | # New function to run YOLO CLI commands 95 | def run_yolo_cli(command_args, capture_output=True, timeout=60): 96 | """ 97 | Run YOLO CLI command and return the results 98 | 99 | Args: 100 | command_args: List of command arguments to pass to yolo CLI 101 | capture_output: Whether to capture and return command output 102 | timeout: Command timeout in seconds 103 | 104 | Returns: 105 | Command output or success status 106 | """ 107 | # Build the complete command 108 | cmd = ["yolo"] + command_args 109 | 110 | # Log the command 111 | logger.info(f"Running YOLO CLI command: {' '.join(cmd)}") 112 | 113 | try: 114 | # Run the command 115 | result = subprocess.run( 116 | cmd, 117 | capture_output=capture_output, 118 | text=True, 119 | check=False, # Don't raise exception on non-zero exit 120 | timeout=timeout 121 | ) 122 | 123 | # Check for errors 124 | if result.returncode != 0: 125 | logger.error(f"YOLO CLI command failed with code {result.returncode}") 126 | logger.error(f"stderr: {result.stderr}") 127 | return { 128 | "success": False, 129 | "error": result.stderr, 130 | "command": " ".join(cmd), 131 | "returncode": result.returncode 132 | } 133 | 134 | # Return the result 135 | if capture_output: 136 | return { 137 | "success": True, 138 | "stdout": result.stdout, 139 | "stderr": result.stderr, 140 | "command": " ".join(cmd) 141 | } 142 | else: 143 | return {"success": True, "command": " ".join(cmd)} 144 | 145 | except subprocess.TimeoutExpired: 146 | logger.error(f"YOLO CLI command timed out after {timeout} seconds") 147 | return { 148 | "success": False, 149 | "error": f"Command timed out after {timeout} seconds", 150 | "command": " ".join(cmd) 151 | } 152 | except Exception as e: 153 | logger.error(f"Error running YOLO CLI command: {str(e)}") 154 | return { 155 | "success": False, 156 | "error": str(e), 157 | "command": " ".join(cmd) 158 | } 159 | 160 | # Create MCP server 161 | mcp = FastMCP("YOLO_Service") 162 | 163 | # Global configuration 164 | CONFIG = { 165 | "model_dirs": [ 166 | ".", # Current directory 167 | "./models", # Models subdirectory 168 | os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"), 169 | ] 170 | } 171 | 172 | # Function to save base64 data to temp file 173 | def save_base64_to_temp(base64_data, prefix="image", suffix=".jpg"): 174 | """Save base64 encoded data to a temporary file and return the path""" 175 | try: 176 | # Create a temporary file 177 | fd, temp_path = tempfile.mkstemp(suffix=suffix, prefix=prefix) 178 | 179 | # Decode base64 data 180 | image_data = base64.b64decode(base64_data) 181 | 182 | # Write data to file 183 | with os.fdopen(fd, 'wb') as temp_file: 184 | temp_file.write(image_data) 185 | 186 | return temp_path 187 | except Exception as e: 188 | logger.error(f"Error saving base64 to temp file: {str(e)}") 189 | raise ValueError(f"Failed to save base64 data: {str(e)}") 190 | 191 | @mcp.tool() 192 | def get_model_directories() -> Dict[str, Any]: 193 | """Get information about configured model directories and available models""" 194 | directories = [] 195 | 196 | for directory in CONFIG["model_dirs"]: 197 | dir_info = { 198 | "path": directory, 199 | "exists": os.path.exists(directory), 200 | "is_directory": os.path.isdir(directory) if os.path.exists(directory) else False, 201 | "models": [] 202 | } 203 | 204 | if dir_info["exists"] and dir_info["is_directory"]: 205 | for filename in os.listdir(directory): 206 | if filename.endswith(".pt"): 207 | dir_info["models"].append(filename) 208 | 209 | directories.append(dir_info) 210 | 211 | return { 212 | "configured_directories": CONFIG["model_dirs"], 213 | "directory_details": directories, 214 | "available_models": list_available_models(), 215 | "loaded_models": [] # No longer track loaded models with CLI approach 216 | } 217 | 218 | @mcp.tool() 219 | def detect_objects( 220 | image_data: str, 221 | model_name: str = "yolov8n.pt", 222 | confidence: float = 0.25, 223 | save_results: bool = False, 224 | is_path: bool = False 225 | ) -> Dict[str, Any]: 226 | """ 227 | Detect objects in an image using YOLO CLI 228 | 229 | Args: 230 | image_data: Base64 encoded image or file path (if is_path=True) 231 | model_name: YOLO model name 232 | confidence: Detection confidence threshold 233 | save_results: Whether to save results to disk 234 | is_path: Whether image_data is a file path 235 | 236 | Returns: 237 | Dictionary containing detection results 238 | """ 239 | try: 240 | # Determine source path 241 | if is_path: 242 | source_path = image_data 243 | if not os.path.exists(source_path): 244 | return { 245 | "error": f"Image file not found: {source_path}", 246 | "source": source_path 247 | } 248 | else: 249 | # Save base64 data to temp file 250 | source_path = save_base64_to_temp(image_data) 251 | 252 | # Determine full model path 253 | model_path = None 254 | for directory in CONFIG["model_dirs"]: 255 | potential_path = os.path.join(directory, model_name) 256 | if os.path.exists(potential_path): 257 | model_path = potential_path 258 | break 259 | 260 | if model_path is None: 261 | available = list_available_models() 262 | available_str = ", ".join(available) if available else "none" 263 | return { 264 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}", 265 | "source": image_data if is_path else "base64_image" 266 | } 267 | 268 | # Setup output directory if saving results 269 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_results") 270 | if save_results and not os.path.exists(output_dir): 271 | os.makedirs(output_dir) 272 | 273 | # Build YOLO CLI command 274 | cmd_args = [ 275 | "detect", # Task 276 | "predict", # Mode 277 | f"model={model_path}", 278 | f"source={source_path}", 279 | f"conf={confidence}", 280 | "format=json", # Request JSON output for parsing 281 | ] 282 | 283 | if save_results: 284 | cmd_args.append(f"project={output_dir}") 285 | cmd_args.append("save=True") 286 | else: 287 | cmd_args.append("save=False") 288 | 289 | # Run YOLO CLI command 290 | result = run_yolo_cli(cmd_args) 291 | 292 | # Clean up temp file if we created one 293 | if not is_path: 294 | try: 295 | os.remove(source_path) 296 | except Exception as e: 297 | logger.warning(f"Failed to clean up temp file {source_path}: {str(e)}") 298 | 299 | # Check for command success 300 | if not result["success"]: 301 | return { 302 | "error": f"YOLO CLI command failed: {result.get('error', 'Unknown error')}", 303 | "command": result.get("command", ""), 304 | "source": image_data if is_path else "base64_image" 305 | } 306 | 307 | # Parse JSON output from stdout 308 | try: 309 | # Try to find JSON in the output 310 | json_start = result["stdout"].find("{") 311 | json_end = result["stdout"].rfind("}") 312 | 313 | if json_start >= 0 and json_end > json_start: 314 | json_str = result["stdout"][json_start:json_end+1] 315 | detection_data = json.loads(json_str) 316 | else: 317 | # If no JSON found, create a basic response with info from stderr 318 | return { 319 | "results": [], 320 | "model_used": model_name, 321 | "total_detections": 0, 322 | "source": image_data if is_path else "base64_image", 323 | "command_output": result["stderr"] 324 | } 325 | 326 | # Format results 327 | formatted_results = [] 328 | 329 | # Parse detection data from YOLO JSON output 330 | if "predictions" in detection_data: 331 | detections = [] 332 | 333 | for pred in detection_data["predictions"]: 334 | # Extract box coordinates 335 | box = pred.get("box", {}) 336 | x1, y1, x2, y2 = box.get("x1", 0), box.get("y1", 0), box.get("x2", 0), box.get("y2", 0) 337 | 338 | # Extract class information 339 | confidence = pred.get("confidence", 0) 340 | class_name = pred.get("name", "unknown") 341 | class_id = pred.get("class", -1) 342 | 343 | detections.append({ 344 | "box": [x1, y1, x2, y2], 345 | "confidence": confidence, 346 | "class_id": class_id, 347 | "class_name": class_name 348 | }) 349 | 350 | # Get image dimensions if available 351 | image_shape = [ 352 | detection_data.get("width", 0), 353 | detection_data.get("height", 0) 354 | ] 355 | 356 | formatted_results.append({ 357 | "detections": detections, 358 | "image_shape": image_shape 359 | }) 360 | 361 | return { 362 | "results": formatted_results, 363 | "model_used": model_name, 364 | "total_detections": sum(len(r["detections"]) for r in formatted_results), 365 | "source": image_data if is_path else "base64_image", 366 | "save_dir": output_dir if save_results else None 367 | } 368 | 369 | except json.JSONDecodeError as e: 370 | logger.error(f"Failed to parse JSON from YOLO output: {e}") 371 | logger.error(f"Output: {result['stdout']}") 372 | 373 | return { 374 | "error": f"Failed to parse YOLO results: {str(e)}", 375 | "command": result.get("command", ""), 376 | "source": image_data if is_path else "base64_image", 377 | "stdout": result.get("stdout", ""), 378 | "stderr": result.get("stderr", "") 379 | } 380 | 381 | except Exception as e: 382 | logger.error(f"Error in detect_objects: {str(e)}") 383 | return { 384 | "error": f"Failed to detect objects: {str(e)}", 385 | "source": image_data if is_path else "base64_image" 386 | } 387 | 388 | @mcp.tool() 389 | def segment_objects( 390 | image_data: str, 391 | model_name: str = "yolov11n-seg.pt", 392 | confidence: float = 0.25, 393 | save_results: bool = False, 394 | is_path: bool = False 395 | ) -> Dict[str, Any]: 396 | """ 397 | Perform instance segmentation on an image using YOLO CLI 398 | 399 | Args: 400 | image_data: Base64 encoded image or file path (if is_path=True) 401 | model_name: YOLO segmentation model name 402 | confidence: Detection confidence threshold 403 | save_results: Whether to save results to disk 404 | is_path: Whether image_data is a file path 405 | 406 | Returns: 407 | Dictionary containing segmentation results 408 | """ 409 | try: 410 | # Determine source path 411 | if is_path: 412 | source_path = image_data 413 | if not os.path.exists(source_path): 414 | return { 415 | "error": f"Image file not found: {source_path}", 416 | "source": source_path 417 | } 418 | else: 419 | # Save base64 data to temp file 420 | source_path = save_base64_to_temp(image_data) 421 | 422 | # Determine full model path 423 | model_path = None 424 | for directory in CONFIG["model_dirs"]: 425 | potential_path = os.path.join(directory, model_name) 426 | if os.path.exists(potential_path): 427 | model_path = potential_path 428 | break 429 | 430 | if model_path is None: 431 | available = list_available_models() 432 | available_str = ", ".join(available) if available else "none" 433 | return { 434 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}", 435 | "source": image_data if is_path else "base64_image" 436 | } 437 | 438 | # Setup output directory if saving results 439 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_results") 440 | if save_results and not os.path.exists(output_dir): 441 | os.makedirs(output_dir) 442 | 443 | # Build YOLO CLI command 444 | cmd_args = [ 445 | "segment", # Task 446 | "predict", # Mode 447 | f"model={model_path}", 448 | f"source={source_path}", 449 | f"conf={confidence}", 450 | "format=json", # Request JSON output for parsing 451 | ] 452 | 453 | if save_results: 454 | cmd_args.append(f"project={output_dir}") 455 | cmd_args.append("save=True") 456 | else: 457 | cmd_args.append("save=False") 458 | 459 | # Run YOLO CLI command 460 | result = run_yolo_cli(cmd_args) 461 | 462 | # Clean up temp file if we created one 463 | if not is_path: 464 | try: 465 | os.remove(source_path) 466 | except Exception as e: 467 | logger.warning(f"Failed to clean up temp file {source_path}: {str(e)}") 468 | 469 | # Check for command success 470 | if not result["success"]: 471 | return { 472 | "error": f"YOLO CLI command failed: {result.get('error', 'Unknown error')}", 473 | "command": result.get("command", ""), 474 | "source": image_data if is_path else "base64_image" 475 | } 476 | 477 | # Parse JSON output from stdout 478 | try: 479 | # Try to find JSON in the output 480 | json_start = result["stdout"].find("{") 481 | json_end = result["stdout"].rfind("}") 482 | 483 | if json_start >= 0 and json_end > json_start: 484 | json_str = result["stdout"][json_start:json_end+1] 485 | segmentation_data = json.loads(json_str) 486 | else: 487 | # If no JSON found, create a basic response with info from stderr 488 | return { 489 | "results": [], 490 | "model_used": model_name, 491 | "total_segments": 0, 492 | "source": image_data if is_path else "base64_image", 493 | "command_output": result["stderr"] 494 | } 495 | 496 | # Format results 497 | formatted_results = [] 498 | 499 | # Parse segmentation data from YOLO JSON output 500 | if "predictions" in segmentation_data: 501 | segments = [] 502 | 503 | for pred in segmentation_data["predictions"]: 504 | # Extract box coordinates 505 | box = pred.get("box", {}) 506 | x1, y1, x2, y2 = box.get("x1", 0), box.get("y1", 0), box.get("x2", 0), box.get("y2", 0) 507 | 508 | # Extract class information 509 | confidence = pred.get("confidence", 0) 510 | class_name = pred.get("name", "unknown") 511 | class_id = pred.get("class", -1) 512 | 513 | segment = { 514 | "box": [x1, y1, x2, y2], 515 | "confidence": confidence, 516 | "class_id": class_id, 517 | "class_name": class_name 518 | } 519 | 520 | # Extract mask if available 521 | if "mask" in pred: 522 | segment["mask"] = pred["mask"] 523 | 524 | segments.append(segment) 525 | 526 | # Get image dimensions if available 527 | image_shape = [ 528 | segmentation_data.get("width", 0), 529 | segmentation_data.get("height", 0) 530 | ] 531 | 532 | formatted_results.append({ 533 | "segments": segments, 534 | "image_shape": image_shape 535 | }) 536 | 537 | return { 538 | "results": formatted_results, 539 | "model_used": model_name, 540 | "total_segments": sum(len(r["segments"]) for r in formatted_results), 541 | "source": image_data if is_path else "base64_image", 542 | "save_dir": output_dir if save_results else None 543 | } 544 | 545 | except json.JSONDecodeError as e: 546 | logger.error(f"Failed to parse JSON from YOLO output: {e}") 547 | logger.error(f"Output: {result['stdout']}") 548 | 549 | return { 550 | "error": f"Failed to parse YOLO results: {str(e)}", 551 | "command": result.get("command", ""), 552 | "source": image_data if is_path else "base64_image", 553 | "stdout": result.get("stdout", ""), 554 | "stderr": result.get("stderr", "") 555 | } 556 | 557 | except Exception as e: 558 | logger.error(f"Error in segment_objects: {str(e)}") 559 | return { 560 | "error": f"Failed to segment objects: {str(e)}", 561 | "source": image_data if is_path else "base64_image" 562 | } 563 | 564 | @mcp.tool() 565 | def classify_image( 566 | image_data: str, 567 | model_name: str = "yolov11n-cls.pt", 568 | top_k: int = 5, 569 | save_results: bool = False, 570 | is_path: bool = False 571 | ) -> Dict[str, Any]: 572 | """ 573 | Classify an image using YOLO classification model via CLI 574 | 575 | Args: 576 | image_data: Base64 encoded image or file path (if is_path=True) 577 | model_name: YOLO classification model name 578 | top_k: Number of top categories to return 579 | save_results: Whether to save results to disk 580 | is_path: Whether image_data is a file path 581 | 582 | Returns: 583 | Dictionary containing classification results 584 | """ 585 | try: 586 | # Determine source path 587 | if is_path: 588 | source_path = image_data 589 | if not os.path.exists(source_path): 590 | return { 591 | "error": f"Image file not found: {source_path}", 592 | "source": source_path 593 | } 594 | else: 595 | # Save base64 data to temp file 596 | source_path = save_base64_to_temp(image_data) 597 | 598 | # Determine full model path 599 | model_path = None 600 | for directory in CONFIG["model_dirs"]: 601 | potential_path = os.path.join(directory, model_name) 602 | if os.path.exists(potential_path): 603 | model_path = potential_path 604 | break 605 | 606 | if model_path is None: 607 | available = list_available_models() 608 | available_str = ", ".join(available) if available else "none" 609 | return { 610 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}", 611 | "source": image_data if is_path else "base64_image" 612 | } 613 | 614 | # Setup output directory if saving results 615 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_results") 616 | if save_results and not os.path.exists(output_dir): 617 | os.makedirs(output_dir) 618 | 619 | # Build YOLO CLI command 620 | cmd_args = [ 621 | "classify", # Task 622 | "predict", # Mode 623 | f"model={model_path}", 624 | f"source={source_path}", 625 | "format=json", # Request JSON output for parsing 626 | ] 627 | 628 | if save_results: 629 | cmd_args.append(f"project={output_dir}") 630 | cmd_args.append("save=True") 631 | else: 632 | cmd_args.append("save=False") 633 | 634 | # Run YOLO CLI command 635 | result = run_yolo_cli(cmd_args) 636 | 637 | # Clean up temp file if we created one 638 | if not is_path: 639 | try: 640 | os.remove(source_path) 641 | except Exception as e: 642 | logger.warning(f"Failed to clean up temp file {source_path}: {str(e)}") 643 | 644 | # Check for command success 645 | if not result["success"]: 646 | return { 647 | "error": f"YOLO CLI command failed: {result.get('error', 'Unknown error')}", 648 | "command": result.get("command", ""), 649 | "source": image_data if is_path else "base64_image" 650 | } 651 | 652 | # Parse JSON output from stdout 653 | try: 654 | # Try to find JSON in the output 655 | json_start = result["stdout"].find("{") 656 | json_end = result["stdout"].rfind("}") 657 | 658 | if json_start >= 0 and json_end > json_start: 659 | json_str = result["stdout"][json_start:json_end+1] 660 | classification_data = json.loads(json_str) 661 | else: 662 | # If no JSON found, create a basic response with info from stderr 663 | return { 664 | "results": [], 665 | "model_used": model_name, 666 | "top_k": top_k, 667 | "source": image_data if is_path else "base64_image", 668 | "command_output": result["stderr"] 669 | } 670 | 671 | # Format results 672 | formatted_results = [] 673 | 674 | # Parse classification data from YOLO JSON output 675 | if "predictions" in classification_data: 676 | classifications = [] 677 | predictions = classification_data["predictions"] 678 | 679 | # Predictions could be an array of classifications 680 | for i, pred in enumerate(predictions[:top_k]): 681 | class_name = pred.get("name", f"class_{i}") 682 | confidence = pred.get("confidence", 0) 683 | 684 | classifications.append({ 685 | "class_id": i, 686 | "class_name": class_name, 687 | "probability": confidence 688 | }) 689 | 690 | # Get image dimensions if available 691 | image_shape = [ 692 | classification_data.get("width", 0), 693 | classification_data.get("height", 0) 694 | ] 695 | 696 | formatted_results.append({ 697 | "classifications": classifications, 698 | "image_shape": image_shape 699 | }) 700 | 701 | return { 702 | "results": formatted_results, 703 | "model_used": model_name, 704 | "top_k": top_k, 705 | "source": image_data if is_path else "base64_image", 706 | "save_dir": output_dir if save_results else None 707 | } 708 | 709 | except json.JSONDecodeError as e: 710 | logger.error(f"Failed to parse JSON from YOLO output: {e}") 711 | logger.error(f"Output: {result['stdout']}") 712 | 713 | return { 714 | "error": f"Failed to parse YOLO results: {str(e)}", 715 | "command": result.get("command", ""), 716 | "source": image_data if is_path else "base64_image", 717 | "stdout": result.get("stdout", ""), 718 | "stderr": result.get("stderr", "") 719 | } 720 | 721 | except Exception as e: 722 | logger.error(f"Error in classify_image: {str(e)}") 723 | return { 724 | "error": f"Failed to classify image: {str(e)}", 725 | "source": image_data if is_path else "base64_image" 726 | } 727 | 728 | @mcp.tool() 729 | def track_objects( 730 | image_data: str, 731 | model_name: str = "yolov8n.pt", 732 | confidence: float = 0.25, 733 | tracker: str = "bytetrack.yaml", 734 | save_results: bool = False 735 | ) -> Dict[str, Any]: 736 | """ 737 | Track objects in an image sequence using YOLO CLI 738 | 739 | Args: 740 | image_data: Base64 encoded image 741 | model_name: YOLO model name 742 | confidence: Detection confidence threshold 743 | tracker: Tracker name to use (e.g., 'bytetrack.yaml', 'botsort.yaml') 744 | save_results: Whether to save results to disk 745 | 746 | Returns: 747 | Dictionary containing tracking results 748 | """ 749 | try: 750 | # Save base64 data to temp file 751 | source_path = save_base64_to_temp(image_data) 752 | 753 | # Determine full model path 754 | model_path = None 755 | for directory in CONFIG["model_dirs"]: 756 | potential_path = os.path.join(directory, model_name) 757 | if os.path.exists(potential_path): 758 | model_path = potential_path 759 | break 760 | 761 | if model_path is None: 762 | available = list_available_models() 763 | available_str = ", ".join(available) if available else "none" 764 | return { 765 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}" 766 | } 767 | 768 | # Setup output directory if saving results 769 | output_dir = os.path.join(tempfile.gettempdir(), "yolo_track_results") 770 | if save_results and not os.path.exists(output_dir): 771 | os.makedirs(output_dir) 772 | 773 | # Build YOLO CLI command 774 | cmd_args = [ 775 | "track", # Combined task and mode for tracking 776 | f"model={model_path}", 777 | f"source={source_path}", 778 | f"conf={confidence}", 779 | f"tracker={tracker}", 780 | "format=json", # Request JSON output for parsing 781 | ] 782 | 783 | if save_results: 784 | cmd_args.append(f"project={output_dir}") 785 | cmd_args.append("save=True") 786 | else: 787 | cmd_args.append("save=False") 788 | 789 | # Run YOLO CLI command 790 | result = run_yolo_cli(cmd_args) 791 | 792 | # Clean up temp file 793 | try: 794 | os.remove(source_path) 795 | except Exception as e: 796 | logger.warning(f"Failed to clean up temp file {source_path}: {str(e)}") 797 | 798 | # Check for command success 799 | if not result["success"]: 800 | return { 801 | "error": f"YOLO CLI command failed: {result.get('error', 'Unknown error')}", 802 | "command": result.get("command", ""), 803 | } 804 | 805 | # Parse JSON output from stdout 806 | try: 807 | # Try to find JSON in the output 808 | json_start = result["stdout"].find("{") 809 | json_end = result["stdout"].rfind("}") 810 | 811 | if json_start >= 0 and json_end > json_start: 812 | json_str = result["stdout"][json_start:json_end+1] 813 | tracking_data = json.loads(json_str) 814 | else: 815 | # If no JSON found, create a basic response 816 | return { 817 | "results": [], 818 | "model_used": model_name, 819 | "tracker": tracker, 820 | "total_tracks": 0, 821 | "command_output": result["stderr"] 822 | } 823 | 824 | # Format results 825 | formatted_results = [] 826 | 827 | # Parse tracking data from YOLO JSON output 828 | if "predictions" in tracking_data: 829 | tracks = [] 830 | 831 | for pred in tracking_data["predictions"]: 832 | # Extract box coordinates 833 | box = pred.get("box", {}) 834 | x1, y1, x2, y2 = box.get("x1", 0), box.get("y1", 0), box.get("x2", 0), box.get("y2", 0) 835 | 836 | # Extract class and tracking information 837 | confidence = pred.get("confidence", 0) 838 | class_name = pred.get("name", "unknown") 839 | class_id = pred.get("class", -1) 840 | track_id = pred.get("id", -1) 841 | 842 | track = { 843 | "box": [x1, y1, x2, y2], 844 | "confidence": confidence, 845 | "class_id": class_id, 846 | "class_name": class_name, 847 | "track_id": track_id 848 | } 849 | 850 | tracks.append(track) 851 | 852 | # Get image dimensions if available 853 | image_shape = [ 854 | tracking_data.get("width", 0), 855 | tracking_data.get("height", 0) 856 | ] 857 | 858 | formatted_results.append({ 859 | "tracks": tracks, 860 | "image_shape": image_shape 861 | }) 862 | 863 | return { 864 | "results": formatted_results, 865 | "model_used": model_name, 866 | "tracker": tracker, 867 | "total_tracks": sum(len(r["tracks"]) for r in formatted_results), 868 | "save_dir": output_dir if save_results else None 869 | } 870 | 871 | except json.JSONDecodeError as e: 872 | logger.error(f"Failed to parse JSON from YOLO output: {e}") 873 | logger.error(f"Output: {result['stdout']}") 874 | 875 | return { 876 | "error": f"Failed to parse YOLO results: {str(e)}", 877 | "command": result.get("command", ""), 878 | "stdout": result.get("stdout", ""), 879 | "stderr": result.get("stderr", "") 880 | } 881 | 882 | except Exception as e: 883 | logger.error(f"Error in track_objects: {str(e)}") 884 | return { 885 | "error": f"Failed to track objects: {str(e)}" 886 | } 887 | 888 | @mcp.tool() 889 | def train_model( 890 | dataset_path: str, 891 | model_name: str = "yolov8n.pt", 892 | epochs: int = 100, 893 | imgsz: int = 640, 894 | batch: int = 16, 895 | name: str = "yolo_custom_model", 896 | project: str = "runs/train" 897 | ) -> Dict[str, Any]: 898 | """ 899 | Train a YOLO model on a custom dataset using CLI 900 | 901 | Args: 902 | dataset_path: Path to YOLO format dataset 903 | model_name: Base model to start with 904 | epochs: Number of training epochs 905 | imgsz: Image size for training 906 | batch: Batch size 907 | name: Name for the training run 908 | project: Project directory 909 | 910 | Returns: 911 | Dictionary containing training results 912 | """ 913 | # Validate dataset path 914 | if not os.path.exists(dataset_path): 915 | return {"error": f"Dataset not found: {dataset_path}"} 916 | 917 | # Determine full model path 918 | model_path = None 919 | for directory in CONFIG["model_dirs"]: 920 | potential_path = os.path.join(directory, model_name) 921 | if os.path.exists(potential_path): 922 | model_path = potential_path 923 | break 924 | 925 | if model_path is None: 926 | available = list_available_models() 927 | available_str = ", ".join(available) if available else "none" 928 | return { 929 | "error": f"Model '{model_name}' not found in any configured directories. Available models: {available_str}" 930 | } 931 | 932 | # Create project directory if it doesn't exist 933 | if not os.path.exists(project): 934 | os.makedirs(project) 935 | 936 | # Determine task type based on model name 937 | task = "detect" # Default task 938 | if "seg" in model_name: 939 | task = "segment" 940 | elif "pose" in model_name: 941 | task = "pose" 942 | elif "cls" in model_name: 943 | task = "classify" 944 | elif "obb" in model_name: 945 | task = "obb" 946 | 947 | # Build YOLO CLI command 948 | cmd_args = [ 949 | task, # Task 950 | "train", # Mode 951 | f"model={model_path}", 952 | f"data={dataset_path}", 953 | f"epochs={epochs}", 954 | f"imgsz={imgsz}", 955 | f"batch={batch}", 956 | f"name={name}", 957 | f"project={project}" 958 | ] 959 | 960 | # Run YOLO CLI command - with longer timeout 961 | logger.info(f"Starting model training with {epochs} epochs - this may take a while...") 962 | result = run_yolo_cli(cmd_args, timeout=epochs * 300) # 5 minutes per epoch 963 | 964 | # Check for command success 965 | if not result["success"]: 966 | return { 967 | "error": f"Training failed: {result.get('error', 'Unknown error')}", 968 | "command": result.get("command", ""), 969 | "stderr": result.get("stderr", "") 970 | } 971 | 972 | # Determine path to best model weights 973 | best_model_path = os.path.join(project, name, "weights", "best.pt") 974 | 975 | # Determine metrics from stdout if possible 976 | metrics = {} 977 | try: 978 | # Look for metrics in output 979 | stdout = result.get("stdout", "") 980 | 981 | # Extract metrics from training output 982 | import re 983 | precision_match = re.search(r"Precision: ([\d\.]+)", stdout) 984 | recall_match = re.search(r"Recall: ([\d\.]+)", stdout) 985 | map50_match = re.search(r"mAP50: ([\d\.]+)", stdout) 986 | map_match = re.search(r"mAP50-95: ([\d\.]+)", stdout) 987 | 988 | if precision_match: 989 | metrics["precision"] = float(precision_match.group(1)) 990 | if recall_match: 991 | metrics["recall"] = float(recall_match.group(1)) 992 | if map50_match: 993 | metrics["mAP50"] = float(map50_match.group(1)) 994 | if map_match: 995 | metrics["mAP50-95"] = float(map_match.group(1)) 996 | except Exception as e: 997 | logger.warning(f"Failed to parse metrics from training output: {str(e)}") 998 | 999 | return { 1000 | "status": "success", 1001 | "model_path": best_model_path, 1002 | "epochs_completed": epochs, 1003 | "final_metrics": metrics, 1004 | "training_log_sample": result.get("stdout", "")[:1000] + "..." if len(result.get("stdout", "")) > 1000 else result.get("stdout", "") 1005 | } 1006 | 1007 | @mcp.tool() 1008 | def validate_model( 1009 | model_path: str, 1010 | data_path: str, 1011 | imgsz: int = 640, 1012 | batch: int = 16 1013 | ) -> Dict[str, Any]: 1014 | """ 1015 | Validate a YOLO model on a dataset using CLI 1016 | 1017 | Args: 1018 | model_path: Path to YOLO model (.pt file) 1019 | data_path: Path to YOLO format validation dataset 1020 | imgsz: Image size for validation 1021 | batch: Batch size 1022 | 1023 | Returns: 1024 | Dictionary containing validation results 1025 | """ 1026 | # Validate model path 1027 | if not os.path.exists(model_path): 1028 | return {"error": f"Model file not found: {model_path}"} 1029 | 1030 | # Validate dataset path 1031 | if not os.path.exists(data_path): 1032 | return {"error": f"Dataset not found: {data_path}"} 1033 | 1034 | # Determine task type based on model name 1035 | model_name = os.path.basename(model_path) 1036 | task = "detect" # Default task 1037 | if "seg" in model_name: 1038 | task = "segment" 1039 | elif "pose" in model_name: 1040 | task = "pose" 1041 | elif "cls" in model_name: 1042 | task = "classify" 1043 | elif "obb" in model_name: 1044 | task = "obb" 1045 | 1046 | # Build YOLO CLI command 1047 | cmd_args = [ 1048 | task, # Task 1049 | "val", # Mode 1050 | f"model={model_path}", 1051 | f"data={data_path}", 1052 | f"imgsz={imgsz}", 1053 | f"batch={batch}" 1054 | ] 1055 | 1056 | # Run YOLO CLI command 1057 | result = run_yolo_cli(cmd_args, timeout=300) # 5 minute timeout 1058 | 1059 | # Check for command success 1060 | if not result["success"]: 1061 | return { 1062 | "error": f"Validation failed: {result.get('error', 'Unknown error')}", 1063 | "command": result.get("command", ""), 1064 | "stderr": result.get("stderr", "") 1065 | } 1066 | 1067 | # Extract metrics from validation output 1068 | metrics = {} 1069 | try: 1070 | stdout = result.get("stdout", "") 1071 | 1072 | import re 1073 | precision_match = re.search(r"Precision: ([\d\.]+)", stdout) 1074 | recall_match = re.search(r"Recall: ([\d\.]+)", stdout) 1075 | map50_match = re.search(r"mAP50: ([\d\.]+)", stdout) 1076 | map_match = re.search(r"mAP50-95: ([\d\.]+)", stdout) 1077 | 1078 | if precision_match: 1079 | metrics["precision"] = float(precision_match.group(1)) 1080 | if recall_match: 1081 | metrics["recall"] = float(recall_match.group(1)) 1082 | if map50_match: 1083 | metrics["mAP50"] = float(map50_match.group(1)) 1084 | if map_match: 1085 | metrics["mAP50-95"] = float(map_match.group(1)) 1086 | except Exception as e: 1087 | logger.warning(f"Failed to parse metrics from validation output: {str(e)}") 1088 | 1089 | return { 1090 | "status": "success", 1091 | "metrics": metrics, 1092 | "validation_output": result.get("stdout", "")[:1000] + "..." if len(result.get("stdout", "")) > 1000 else result.get("stdout", "") 1093 | } 1094 | 1095 | @mcp.tool() 1096 | def export_model( 1097 | model_path: str, 1098 | format: str = "onnx", 1099 | imgsz: int = 640 1100 | ) -> Dict[str, Any]: 1101 | """ 1102 | Export a YOLO model to different formats using CLI 1103 | 1104 | Args: 1105 | model_path: Path to YOLO model (.pt file) 1106 | format: Export format (onnx, torchscript, openvino, etc.) 1107 | imgsz: Image size for export 1108 | 1109 | Returns: 1110 | Dictionary containing export results 1111 | """ 1112 | # Validate model path 1113 | if not os.path.exists(model_path): 1114 | return {"error": f"Model file not found: {model_path}"} 1115 | 1116 | # Valid export formats 1117 | valid_formats = [ 1118 | "torchscript", "onnx", "openvino", "engine", "coreml", "saved_model", 1119 | "pb", "tflite", "edgetpu", "tfjs", "paddle" 1120 | ] 1121 | 1122 | if format not in valid_formats: 1123 | return {"error": f"Invalid export format: {format}. Valid formats include: {', '.join(valid_formats)}"} 1124 | 1125 | # Build YOLO CLI command 1126 | cmd_args = [ 1127 | "export", # Combined task and mode for export 1128 | f"model={model_path}", 1129 | f"format={format}", 1130 | f"imgsz={imgsz}" 1131 | ] 1132 | 1133 | # Run YOLO CLI command 1134 | result = run_yolo_cli(cmd_args, timeout=300) # 5 minute timeout 1135 | 1136 | # Check for command success 1137 | if not result["success"]: 1138 | return { 1139 | "error": f"Export failed: {result.get('error', 'Unknown error')}", 1140 | "command": result.get("command", ""), 1141 | "stderr": result.get("stderr", "") 1142 | } 1143 | 1144 | # Try to determine export path 1145 | export_path = None 1146 | try: 1147 | # Model path without extension 1148 | base_path = os.path.splitext(model_path)[0] 1149 | 1150 | # Expected export paths based on format 1151 | format_extensions = { 1152 | "torchscript": ".torchscript", 1153 | "onnx": ".onnx", 1154 | "openvino": "_openvino_model", 1155 | "engine": ".engine", 1156 | "coreml": ".mlmodel", 1157 | "saved_model": "_saved_model", 1158 | "pb": ".pb", 1159 | "tflite": ".tflite", 1160 | "edgetpu": "_edgetpu.tflite", 1161 | "tfjs": "_web_model", 1162 | "paddle": "_paddle_model" 1163 | } 1164 | 1165 | expected_ext = format_extensions.get(format, "") 1166 | expected_path = base_path + expected_ext 1167 | 1168 | # Check if the exported file exists 1169 | if os.path.exists(expected_path) or os.path.isdir(expected_path): 1170 | export_path = expected_path 1171 | except Exception as e: 1172 | logger.warning(f"Failed to determine export path: {str(e)}") 1173 | 1174 | return { 1175 | "status": "success", 1176 | "export_path": export_path, 1177 | "format": format, 1178 | "export_output": result.get("stdout", "")[:1000] + "..." if len(result.get("stdout", "")) > 1000 else result.get("stdout", "") 1179 | } 1180 | 1181 | @mcp.tool() 1182 | def list_available_models() -> List[str]: 1183 | """List available YOLO models that actually exist on disk in any configured directory""" 1184 | # Common YOLO model patterns 1185 | model_patterns = [ 1186 | "yolov11*.pt", 1187 | "yolov8*.pt" 1188 | ] 1189 | 1190 | # Find all existing models in all configured directories 1191 | available_models = set() 1192 | for directory in CONFIG["model_dirs"]: 1193 | if not os.path.exists(directory): 1194 | continue 1195 | 1196 | # Check for model files directly 1197 | for filename in os.listdir(directory): 1198 | if filename.endswith(".pt") and any( 1199 | fnmatch.fnmatch(filename, pattern) for pattern in model_patterns 1200 | ): 1201 | available_models.add(filename) 1202 | 1203 | # Convert to sorted list 1204 | result = sorted(list(available_models)) 1205 | 1206 | if not result: 1207 | logger.warning("No model files found in configured directories.") 1208 | return ["No models available - download models to any of these directories: " + ", ".join(CONFIG["model_dirs"])] 1209 | 1210 | return result 1211 | 1212 | # Camera detection functions using CLI instead of Python API 1213 | def camera_detection_thread(model_name, confidence, fps_limit=30, camera_id=0): 1214 | """Background thread for camera detection using YOLO CLI""" 1215 | global camera_running, detection_results, camera_last_access_time, camera_startup_status, camera_last_error 1216 | 1217 | try: 1218 | # Create a unique directory for camera results 1219 | output_dir = os.path.join(tempfile.gettempdir(), f"yolo_camera_{int(time.time())}") 1220 | os.makedirs(output_dir, exist_ok=True) 1221 | 1222 | # Determine full model path 1223 | model_path = None 1224 | for directory in CONFIG["model_dirs"]: 1225 | potential_path = os.path.join(directory, model_name) 1226 | if os.path.exists(potential_path): 1227 | model_path = potential_path 1228 | break 1229 | 1230 | if model_path is None: 1231 | error_msg = f"Model {model_name} not found in any configured directories" 1232 | logger.error(error_msg) 1233 | camera_running = False 1234 | camera_startup_status = { 1235 | "success": False, 1236 | "error": error_msg, 1237 | "timestamp": time.time() 1238 | } 1239 | detection_results.append({ 1240 | "timestamp": time.time(), 1241 | "error": f"Failed to load model: Model not found", 1242 | "camera_status": "error", 1243 | "detections": [] 1244 | }) 1245 | return 1246 | 1247 | # Log camera start 1248 | logger.info(f"Starting camera detection with model {model_name}, camera ID {camera_id}") 1249 | detection_results.append({ 1250 | "timestamp": time.time(), 1251 | "system_info": { 1252 | "os": platform.system() if 'platform' in globals() else "Unknown", 1253 | "camera_id": camera_id 1254 | }, 1255 | "camera_status": "starting", 1256 | "detections": [] 1257 | }) 1258 | 1259 | # Determine task type based on model name 1260 | task = "detect" # Default task 1261 | if "seg" in model_name: 1262 | task = "segment" 1263 | elif "pose" in model_name: 1264 | task = "pose" 1265 | elif "cls" in model_name: 1266 | task = "classify" 1267 | 1268 | # Build YOLO CLI command 1269 | base_cmd_args = [ 1270 | task, # Task 1271 | "predict", # Mode 1272 | f"model={model_path}", 1273 | f"source={camera_id}", # Camera source ID 1274 | f"conf={confidence}", 1275 | "format=json", 1276 | "save=False", # Don't save frames by default 1277 | "show=False" # Don't show GUI window 1278 | ] 1279 | 1280 | # First verify YOLO command is available 1281 | logger.info("Verifying YOLO CLI availability before starting camera...") 1282 | check_cmd = ["yolo", "--version"] 1283 | try: 1284 | check_result = subprocess.run( 1285 | check_cmd, 1286 | capture_output=True, 1287 | text=True, 1288 | check=False, 1289 | timeout=10 1290 | ) 1291 | 1292 | if check_result.returncode != 0: 1293 | error_msg = f"YOLO CLI check failed with code {check_result.returncode}: {check_result.stderr}" 1294 | logger.error(error_msg) 1295 | camera_running = False 1296 | camera_startup_status = { 1297 | "success": False, 1298 | "error": error_msg, 1299 | "timestamp": time.time() 1300 | } 1301 | detection_results.append({ 1302 | "timestamp": time.time(), 1303 | "error": error_msg, 1304 | "camera_status": "error", 1305 | "detections": [] 1306 | }) 1307 | return 1308 | 1309 | logger.info(f"YOLO CLI is available: {check_result.stdout.strip()}") 1310 | except Exception as e: 1311 | error_msg = f"Error checking YOLO CLI: {str(e)}" 1312 | logger.error(error_msg) 1313 | camera_running = False 1314 | camera_startup_status = { 1315 | "success": False, 1316 | "error": error_msg, 1317 | "timestamp": time.time() 1318 | } 1319 | detection_results.append({ 1320 | "timestamp": time.time(), 1321 | "error": error_msg, 1322 | "camera_status": "error", 1323 | "detections": [] 1324 | }) 1325 | return 1326 | 1327 | # Set up subprocess for ongoing camera capture 1328 | process = None 1329 | frame_count = 0 1330 | error_count = 0 1331 | start_time = time.time() 1332 | 1333 | # Start YOLO CLI process 1334 | cmd_str = "yolo " + " ".join(base_cmd_args) 1335 | logger.info(f"Starting YOLO CLI process: {cmd_str}") 1336 | 1337 | try: 1338 | process = subprocess.Popen( 1339 | ["yolo"] + base_cmd_args, 1340 | stdin=subprocess.PIPE, 1341 | stdout=subprocess.PIPE, 1342 | stderr=subprocess.PIPE, 1343 | text=True, 1344 | bufsize=1, # Line buffered 1345 | ) 1346 | 1347 | # Wait a moment to check if the process immediately fails 1348 | time.sleep(1) 1349 | if process.poll() is not None: 1350 | error_msg = f"YOLO process failed to start (exit code {process.returncode})" 1351 | stderr_output = process.stderr.read() 1352 | logger.error(f"{error_msg} - STDERR: {stderr_output}") 1353 | 1354 | camera_running = False 1355 | camera_startup_status = { 1356 | "success": False, 1357 | "error": error_msg, 1358 | "stderr": stderr_output, 1359 | "timestamp": time.time() 1360 | } 1361 | detection_results.append({ 1362 | "timestamp": time.time(), 1363 | "error": error_msg, 1364 | "stderr": stderr_output, 1365 | "camera_status": "error", 1366 | "detections": [] 1367 | }) 1368 | return 1369 | 1370 | # Process started successfully 1371 | camera_startup_status = { 1372 | "success": True, 1373 | "timestamp": time.time() 1374 | } 1375 | 1376 | # Handle camera stream 1377 | while camera_running: 1378 | # Read output line from process 1379 | stdout_line = process.stdout.readline().strip() 1380 | 1381 | if not stdout_line: 1382 | # Check if process is still running 1383 | if process.poll() is not None: 1384 | error_msg = f"YOLO process ended unexpectedly with code {process.returncode}" 1385 | stderr_output = process.stderr.read() 1386 | logger.error(f"{error_msg} - STDERR: {stderr_output}") 1387 | 1388 | camera_running = False 1389 | camera_last_error = { 1390 | "error": error_msg, 1391 | "stderr": stderr_output, 1392 | "timestamp": time.time() 1393 | } 1394 | detection_results.append({ 1395 | "timestamp": time.time(), 1396 | "error": error_msg, 1397 | "camera_status": "error", 1398 | "stderr": stderr_output, 1399 | "detections": [] 1400 | }) 1401 | break 1402 | 1403 | time.sleep(0.1) # Short sleep to avoid CPU spin 1404 | continue 1405 | 1406 | # Try to parse JSON output from YOLO 1407 | try: 1408 | # Find JSON in the output line 1409 | json_start = stdout_line.find("{") 1410 | if json_start >= 0: 1411 | json_str = stdout_line[json_start:] 1412 | detection_data = json.loads(json_str) 1413 | 1414 | frame_count += 1 1415 | 1416 | # Process detection data 1417 | if "predictions" in detection_data: 1418 | detections = [] 1419 | 1420 | for pred in detection_data["predictions"]: 1421 | # Extract box coordinates 1422 | box = pred.get("box", {}) 1423 | x1, y1, x2, y2 = box.get("x1", 0), box.get("y1", 0), box.get("x2", 0), box.get("y2", 0) 1424 | 1425 | # Extract class information 1426 | confidence = pred.get("confidence", 0) 1427 | class_name = pred.get("name", "unknown") 1428 | class_id = pred.get("class", -1) 1429 | 1430 | detections.append({ 1431 | "box": [x1, y1, x2, y2], 1432 | "confidence": confidence, 1433 | "class_id": class_id, 1434 | "class_name": class_name 1435 | }) 1436 | 1437 | # Update detection results (keep only the last 10) 1438 | if len(detection_results) >= 10: 1439 | detection_results.pop(0) 1440 | 1441 | # Get image dimensions if available 1442 | image_shape = [ 1443 | detection_data.get("width", 0), 1444 | detection_data.get("height", 0) 1445 | ] 1446 | 1447 | detection_results.append({ 1448 | "timestamp": time.time(), 1449 | "frame_count": frame_count, 1450 | "detections": detections, 1451 | "camera_status": "running", 1452 | "image_shape": image_shape 1453 | }) 1454 | 1455 | # Update last access time when processing frames 1456 | camera_last_access_time = time.time() 1457 | 1458 | # Log occasional status 1459 | if frame_count % 30 == 0: 1460 | fps = frame_count / (time.time() - start_time) 1461 | logger.info(f"Camera running: processed {frame_count} frames ({fps:.1f} FPS)") 1462 | detection_count = sum(len(r.get("detections", [])) for r in detection_results if "detections" in r) 1463 | logger.info(f"Total detections in current buffer: {detection_count}") 1464 | 1465 | except json.JSONDecodeError: 1466 | # Not all lines will be valid JSON, that's normal 1467 | pass 1468 | except Exception as e: 1469 | error_msg = f"Error processing camera output: {str(e)}" 1470 | logger.warning(error_msg) 1471 | error_count += 1 1472 | 1473 | if error_count > 10: 1474 | logger.error("Too many processing errors, stopping camera") 1475 | camera_running = False 1476 | camera_last_error = { 1477 | "error": "Too many processing errors", 1478 | "timestamp": time.time() 1479 | } 1480 | break 1481 | 1482 | except Exception as e: 1483 | error_msg = f"Error in camera process management: {str(e)}" 1484 | logger.error(error_msg) 1485 | camera_running = False 1486 | camera_startup_status = { 1487 | "success": False, 1488 | "error": error_msg, 1489 | "timestamp": time.time() 1490 | } 1491 | detection_results.append({ 1492 | "timestamp": time.time(), 1493 | "error": error_msg, 1494 | "camera_status": "error", 1495 | "detections": [] 1496 | }) 1497 | return 1498 | 1499 | except Exception as e: 1500 | error_msg = f"Error in camera thread: {str(e)}" 1501 | logger.error(error_msg) 1502 | camera_running = False 1503 | camera_startup_status = { 1504 | "success": False, 1505 | "error": error_msg, 1506 | "timestamp": time.time() 1507 | } 1508 | detection_results.append({ 1509 | "timestamp": time.time(), 1510 | "error": error_msg, 1511 | "camera_status": "error", 1512 | "detections": [] 1513 | }) 1514 | 1515 | finally: 1516 | # Clean up 1517 | logger.info("Shutting down camera...") 1518 | camera_running = False 1519 | 1520 | if process is not None and process.poll() is None: 1521 | try: 1522 | # Terminate process 1523 | process.terminate() 1524 | process.wait(timeout=5) 1525 | except subprocess.TimeoutExpired: 1526 | process.kill() # Force kill if terminate doesn't work 1527 | except Exception as e: 1528 | logger.error(f"Error terminating YOLO process: {str(e)}") 1529 | 1530 | logger.info("Camera detection stopped") 1531 | 1532 | 1533 | @mcp.tool() 1534 | def start_camera_detection( 1535 | model_name: str = "yolov8n.pt", 1536 | confidence: float = 0.25, 1537 | camera_id: int = 0 1538 | ) -> Dict[str, Any]: 1539 | """ 1540 | Start realtime object detection using the computer's camera via YOLO CLI 1541 | 1542 | Args: 1543 | model_name: YOLO model name to use 1544 | confidence: Detection confidence threshold 1545 | camera_id: Camera device ID (0 is usually the default camera) 1546 | 1547 | Returns: 1548 | Status of camera detection 1549 | """ 1550 | global camera_thread, camera_running, detection_results, camera_last_access_time, camera_startup_status, camera_last_error 1551 | 1552 | # Reset status variables 1553 | camera_startup_status = None 1554 | camera_last_error = None 1555 | 1556 | # Check if already running 1557 | if camera_running: 1558 | # Update last access time 1559 | camera_last_access_time = time.time() 1560 | return {"status": "success", "message": "Camera detection is already running"} 1561 | 1562 | # Clear previous results 1563 | detection_results = [] 1564 | 1565 | # First check if YOLO CLI is available 1566 | try: 1567 | version_check = run_yolo_cli(["--version"], timeout=10) 1568 | if not version_check["success"]: 1569 | return { 1570 | "status": "error", 1571 | "message": "YOLO CLI not available or not properly installed", 1572 | "details": version_check.get("error", "Unknown error"), 1573 | "solution": "Please make sure the 'yolo' command is in your PATH" 1574 | } 1575 | except Exception as e: 1576 | error_msg = f"Error checking YOLO CLI: {str(e)}" 1577 | logger.error(error_msg) 1578 | return { 1579 | "status": "error", 1580 | "message": error_msg, 1581 | "solution": "Please make sure the 'yolo' command is in your PATH" 1582 | } 1583 | 1584 | # Start detection thread 1585 | camera_running = True 1586 | camera_last_access_time = time.time() # Update access time 1587 | camera_thread = threading.Thread( 1588 | target=camera_detection_thread, 1589 | args=(model_name, confidence, 30, camera_id), 1590 | daemon=True 1591 | ) 1592 | camera_thread.start() 1593 | 1594 | # Give the thread a moment to initialize and potentially fail 1595 | time.sleep(2) 1596 | 1597 | # Check if the thread has reported any startup issues 1598 | if camera_startup_status and not camera_startup_status.get("success", False): 1599 | # Camera thread encountered an error during startup 1600 | return { 1601 | "status": "error", 1602 | "message": "Camera detection failed to start", 1603 | "details": camera_startup_status, 1604 | "solution": "Check logs for detailed error information" 1605 | } 1606 | 1607 | # Thread is running, camera should be starting 1608 | return { 1609 | "status": "success", 1610 | "message": f"Started camera detection using model {model_name}", 1611 | "model": model_name, 1612 | "confidence": confidence, 1613 | "camera_id": camera_id, 1614 | "auto_shutdown": f"Camera will auto-shutdown after {CAMERA_INACTIVITY_TIMEOUT} seconds of inactivity", 1615 | "note": "If camera doesn't work, try different camera_id values (0, 1, or 2)" 1616 | } 1617 | 1618 | 1619 | @mcp.tool() 1620 | def stop_camera_detection() -> Dict[str, Any]: 1621 | """ 1622 | Stop realtime camera detection 1623 | 1624 | Returns: 1625 | Status message 1626 | """ 1627 | global camera_running 1628 | 1629 | if not camera_running: 1630 | return {"status": "error", "message": "Camera detection is not running"} 1631 | 1632 | logger.info("Stopping camera detection by user request") 1633 | camera_running = False 1634 | 1635 | # Wait for thread to terminate 1636 | if camera_thread and camera_thread.is_alive(): 1637 | camera_thread.join(timeout=2.0) 1638 | 1639 | return { 1640 | "status": "success", 1641 | "message": "Stopped camera detection" 1642 | } 1643 | 1644 | @mcp.tool() 1645 | def get_camera_detections() -> Dict[str, Any]: 1646 | """ 1647 | Get the latest detections from the camera 1648 | 1649 | Returns: 1650 | Dictionary with recent detections 1651 | """ 1652 | global detection_results, camera_thread, camera_last_access_time, camera_startup_status, camera_last_error 1653 | 1654 | # Update the last access time whenever this function is called 1655 | if camera_running: 1656 | camera_last_access_time = time.time() 1657 | 1658 | # Check if thread is alive 1659 | thread_alive = camera_thread is not None and camera_thread.is_alive() 1660 | 1661 | # If camera_running is False, check if we have startup status information 1662 | if not camera_running and camera_startup_status and not camera_startup_status.get("success", False): 1663 | return { 1664 | "status": "error", 1665 | "message": "Camera detection failed to start", 1666 | "is_running": False, 1667 | "camera_status": "error", 1668 | "startup_error": camera_startup_status, 1669 | "solution": "Check logs for detailed error information" 1670 | } 1671 | 1672 | # If camera_running is True but thread is dead, there's an issue 1673 | if camera_running and not thread_alive: 1674 | return { 1675 | "status": "error", 1676 | "message": "Camera thread has stopped unexpectedly", 1677 | "is_running": False, 1678 | "camera_status": "error", 1679 | "thread_alive": thread_alive, 1680 | "last_error": camera_last_error, 1681 | "detections": detection_results, 1682 | "count": len(detection_results), 1683 | "solution": "Please try restart the camera with a different camera_id" 1684 | } 1685 | 1686 | if not camera_running: 1687 | return { 1688 | "status": "error", 1689 | "message": "Camera detection is not running", 1690 | "is_running": False, 1691 | "camera_status": "stopped" 1692 | } 1693 | 1694 | # Check for errors in detection results 1695 | errors = [result.get("error") for result in detection_results if "error" in result] 1696 | recent_errors = errors[-5:] if errors else [] 1697 | 1698 | # Count actual detections 1699 | detection_count = sum(len(result.get("detections", [])) for result in detection_results if "detections" in result) 1700 | 1701 | return { 1702 | "status": "success", 1703 | "is_running": camera_running, 1704 | "thread_alive": thread_alive, 1705 | "detections": detection_results, 1706 | "count": len(detection_results), 1707 | "total_detections": detection_count, 1708 | "recent_errors": recent_errors if recent_errors else None, 1709 | "camera_status": "error" if recent_errors else "running", 1710 | "inactivity_timeout": { 1711 | "seconds_remaining": int(CAMERA_INACTIVITY_TIMEOUT - (time.time() - camera_last_access_time)), 1712 | "last_access": camera_last_access_time 1713 | } 1714 | } 1715 | 1716 | 1717 | @mcp.tool() 1718 | def comprehensive_image_analysis( 1719 | image_path: str, 1720 | confidence: float = 0.25, 1721 | save_results: bool = False 1722 | ) -> Dict[str, Any]: 1723 | """ 1724 | Perform comprehensive analysis on an image by combining multiple CLI model results 1725 | 1726 | Args: 1727 | image_path: Path to the image file 1728 | confidence: Detection confidence threshold 1729 | save_results: Whether to save results to disk 1730 | 1731 | Returns: 1732 | Dictionary containing comprehensive analysis results 1733 | """ 1734 | try: 1735 | if not os.path.exists(image_path): 1736 | return {"error": f"Image file not found: {image_path}"} 1737 | 1738 | analysis_results = {} 1739 | 1740 | # 1. Object detection 1741 | logger.info("Running object detection for comprehensive analysis") 1742 | object_result = detect_objects( 1743 | image_data=image_path, 1744 | model_name="yolov11n.pt", 1745 | confidence=confidence, 1746 | save_results=save_results, 1747 | is_path=True 1748 | ) 1749 | 1750 | # Process object detection results 1751 | detected_objects = [] 1752 | if "results" in object_result and object_result["results"]: 1753 | for result in object_result["results"]: 1754 | for obj in result.get("detections", []): 1755 | detected_objects.append({ 1756 | "class_name": obj.get("class_name", "unknown"), 1757 | "confidence": obj.get("confidence", 0) 1758 | }) 1759 | analysis_results["objects"] = detected_objects 1760 | 1761 | # 2. Scene classification 1762 | try: 1763 | logger.info("Running classification for comprehensive analysis") 1764 | cls_result = classify_image( 1765 | image_data=image_path, 1766 | model_name="yolov8n-cls.pt", 1767 | top_k=3, 1768 | save_results=False, 1769 | is_path=True 1770 | ) 1771 | 1772 | scene_classifications = [] 1773 | if "results" in cls_result and cls_result["results"]: 1774 | for result in cls_result["results"]: 1775 | for cls in result.get("classifications", []): 1776 | scene_classifications.append({ 1777 | "class_name": cls.get("class_name", "unknown"), 1778 | "probability": cls.get("probability", 0) 1779 | }) 1780 | analysis_results["scene"] = scene_classifications 1781 | except Exception as e: 1782 | logger.error(f"Error during scene classification: {str(e)}") 1783 | analysis_results["scene_error"] = str(e) 1784 | 1785 | # 3. Human pose detection (if pose model is available) 1786 | try: 1787 | # Check if pose model exists 1788 | pose_model_exists = False 1789 | for directory in CONFIG["model_dirs"]: 1790 | if os.path.exists(os.path.join(directory, "yolov8n-pose.pt")): 1791 | pose_model_exists = True 1792 | break 1793 | 1794 | if pose_model_exists: 1795 | logger.info("Running pose detection for comprehensive analysis") 1796 | # Build YOLO CLI command for pose detection 1797 | cmd_args = [ 1798 | "pose", # Task 1799 | "predict", # Mode 1800 | f"model=yolov8n-pose.pt", 1801 | f"source={image_path}", 1802 | f"conf={confidence}", 1803 | "format=json", 1804 | ] 1805 | 1806 | result = run_yolo_cli(cmd_args) 1807 | 1808 | if result["success"]: 1809 | # Parse JSON output 1810 | json_start = result["stdout"].find("{") 1811 | json_end = result["stdout"].rfind("}") 1812 | 1813 | if json_start >= 0 and json_end > json_start: 1814 | json_str = result["stdout"][json_start:json_end+1] 1815 | pose_data = json.loads(json_str) 1816 | 1817 | detected_poses = [] 1818 | if "predictions" in pose_data: 1819 | for pred in pose_data["predictions"]: 1820 | confidence = pred.get("confidence", 0) 1821 | keypoints = pred.get("keypoints", []) 1822 | 1823 | detected_poses.append({ 1824 | "person_confidence": confidence, 1825 | "has_keypoints": len(keypoints) if keypoints else 0 1826 | }) 1827 | 1828 | analysis_results["poses"] = detected_poses 1829 | else: 1830 | analysis_results["pose_error"] = "Pose model not available" 1831 | 1832 | except Exception as e: 1833 | logger.error(f"Error during pose detection: {str(e)}") 1834 | analysis_results["pose_error"] = str(e) 1835 | 1836 | # 4. Comprehensive task description 1837 | tasks = [] 1838 | 1839 | # Detect main objects 1840 | main_objects = [obj["class_name"] for obj in detected_objects if obj["confidence"] > 0.5] 1841 | if "person" in main_objects: 1842 | tasks.append("Person Detection") 1843 | 1844 | # Check for weapon objects 1845 | weapon_objects = ["sword", "knife", "katana", "gun", "pistol", "rifle"] 1846 | weapons = [obj for obj in main_objects if any(weapon in obj.lower() for weapon in weapon_objects)] 1847 | if weapons: 1848 | tasks.append(f"Weapon Detection ({', '.join(weapons)})") 1849 | 1850 | # Count people 1851 | person_count = main_objects.count("person") 1852 | if person_count > 0: 1853 | tasks.append(f"Person Count ({person_count} people)") 1854 | 1855 | # Pose analysis 1856 | if "poses" in analysis_results and analysis_results["poses"]: 1857 | tasks.append("Human Pose Analysis") 1858 | 1859 | # Scene classification 1860 | if "scene" in analysis_results and analysis_results["scene"]: 1861 | scene_types = [scene["class_name"] for scene in analysis_results["scene"][:2]] 1862 | tasks.append(f"Scene Classification ({', '.join(scene_types)})") 1863 | 1864 | analysis_results["identified_tasks"] = tasks 1865 | 1866 | # Return comprehensive results 1867 | return { 1868 | "status": "success", 1869 | "image_path": image_path, 1870 | "analysis": analysis_results, 1871 | "summary": "Tasks identified in the image: " + ", ".join(tasks) if tasks else "No clear tasks identified" 1872 | } 1873 | except Exception as e: 1874 | return { 1875 | "status": "error", 1876 | "image_path": image_path, 1877 | "error": f"Comprehensive analysis failed: {str(e)}" 1878 | } 1879 | 1880 | @mcp.tool() 1881 | def analyze_image_from_path( 1882 | image_path: str, 1883 | model_name: str = "yolov8n.pt", 1884 | confidence: float = 0.25, 1885 | save_results: bool = False 1886 | ) -> Dict[str, Any]: 1887 | """ 1888 | Analyze image from file path using YOLO CLI 1889 | 1890 | Args: 1891 | image_path: Path to the image file 1892 | model_name: YOLO model name 1893 | confidence: Detection confidence threshold 1894 | save_results: Whether to save results to disk 1895 | 1896 | Returns: 1897 | Dictionary containing detection results 1898 | """ 1899 | try: 1900 | # Call detect_objects function with is_path=True 1901 | return detect_objects( 1902 | image_data=image_path, 1903 | model_name=model_name, 1904 | confidence=confidence, 1905 | save_results=save_results, 1906 | is_path=True 1907 | ) 1908 | except Exception as e: 1909 | return { 1910 | "error": f"Failed to analyze image: {str(e)}", 1911 | "image_path": image_path 1912 | } 1913 | 1914 | @mcp.tool() 1915 | def test_connection() -> Dict[str, Any]: 1916 | """ 1917 | Test if YOLO CLI service is running properly 1918 | 1919 | Returns: 1920 | Status information and available tools 1921 | """ 1922 | # Test YOLO CLI availability 1923 | try: 1924 | version_result = run_yolo_cli(["--version"], timeout=10) 1925 | yolo_version = version_result.get("stdout", "Unknown") if version_result.get("success") else "Not available" 1926 | 1927 | # Clean up version string 1928 | if "ultralytics" in yolo_version.lower(): 1929 | yolo_version = yolo_version.strip() 1930 | else: 1931 | yolo_version = "YOLO CLI not found or not responding correctly" 1932 | except Exception as e: 1933 | yolo_version = f"Error checking YOLO CLI: {str(e)}" 1934 | 1935 | return { 1936 | "status": "YOLO CLI service is running normally", 1937 | "yolo_version": yolo_version, 1938 | "available_models": list_available_models(), 1939 | "available_tools": [ 1940 | "list_available_models", "detect_objects", "segment_objects", 1941 | "classify_image", "track_objects", "train_model", "validate_model", 1942 | "export_model", "start_camera_detection", "stop_camera_detection", 1943 | "get_camera_detections", "test_connection", 1944 | # Additional tools 1945 | "analyze_image_from_path", 1946 | "comprehensive_image_analysis" 1947 | ], 1948 | "features": [ 1949 | "All detection functions use YOLO CLI rather than Python API", 1950 | "Support for loading images directly from file paths", 1951 | "Support for comprehensive image analysis with task identification", 1952 | "Support for camera detection using YOLO CLI" 1953 | ] 1954 | } 1955 | 1956 | def cleanup_resources(): 1957 | """Clean up resources when the server is shutting down""" 1958 | global camera_running 1959 | 1960 | logger.info("Cleaning up resources...") 1961 | 1962 | # Stop camera if it's running 1963 | if camera_running: 1964 | logger.info("Shutting down camera during server exit") 1965 | camera_running = False 1966 | 1967 | # Give the camera thread a moment to clean up 1968 | if camera_thread and camera_thread.is_alive(): 1969 | camera_thread.join(timeout=2.0) 1970 | 1971 | logger.info("Cleanup complete") 1972 | 1973 | def signal_handler(sig, frame): 1974 | """Handle termination signals""" 1975 | logger.info(f"Received signal {sig}, shutting down...") 1976 | cleanup_resources() 1977 | sys.exit(0) 1978 | 1979 | def start_watchdog(): 1980 | """Start the camera watchdog thread""" 1981 | watchdog = threading.Thread( 1982 | target=camera_watchdog_thread, 1983 | daemon=True 1984 | ) 1985 | watchdog.start() 1986 | return watchdog 1987 | 1988 | # Register cleanup functions 1989 | atexit.register(cleanup_resources) 1990 | signal.signal(signal.SIGINT, signal_handler) 1991 | signal.signal(signal.SIGTERM, signal_handler) 1992 | 1993 | # Modify the main execution section 1994 | if __name__ == "__main__": 1995 | import platform 1996 | 1997 | logger.info("Starting YOLO CLI service") 1998 | logger.info(f"Platform: {platform.system()} {platform.release()}") 1999 | 2000 | # Test if YOLO CLI is available 2001 | try: 2002 | test_result = run_yolo_cli(["--version"], timeout=10) 2003 | if test_result["success"]: 2004 | logger.info(f"YOLO CLI available: {test_result.get('stdout', '').strip()}") 2005 | else: 2006 | logger.warning(f"YOLO CLI test failed: {test_result.get('stderr', '')}") 2007 | logger.warning("Service may not function correctly without YOLO CLI available") 2008 | except Exception as e: 2009 | logger.error(f"Error testing YOLO CLI: {str(e)}") 2010 | logger.warning("Service may not function correctly without YOLO CLI available") 2011 | 2012 | # Start the camera watchdog thread 2013 | watchdog_thread = start_watchdog() 2014 | 2015 | # Initialize and run server 2016 | logger.info("Starting MCP server...") 2017 | mcp.run(transport='stdio') --------------------------------------------------------------------------------