ship-detector-mcp 0.1.0__tar.gz

ship_detector_mcp-0.1.0/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: ship-detector-mcp
+Version: 0.1.0
+Summary: A MCP server for ship detection using YOLO
+Author-email: xiacb <chenboxia111@163.com>
+Requires-Python: >=3.10
+Requires-Dist: mcp[cli]>=1.9.2
+Requires-Dist: numpy
+Requires-Dist: opencv-python
+Requires-Dist: pillow
+Requires-Dist: requests
+Requires-Dist: ultralytics

ship_detector_mcp-0.1.0/pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "ship-detector-mcp"
+version = "0.1.0"
+description = "A MCP server for ship detection using YOLO"
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [
+    {name = "xiacb", email = "chenboxia111@163.com"}
+]
+dependencies = [
+    "mcp[cli]>=1.9.2",
+    "ultralytics",
+    "opencv-python",
+    "pillow",
+    "numpy",
+    "requests",
+]
+
+[project.scripts]
+ship-detector-mcp = "ship_detector_mcp:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/ship_detector_mcp"]

ship_detector_mcp-0.1.0/src/ship_detector_mcp/__init__.py

@@ -0,0 +1,11 @@
+import argparse
+from .server import mcp
+
+def main():
+    """启动 MCP 服务"""
+    parser = argparse.ArgumentParser(description="Ship Detector MCP Server")
+    parser.parse_args()
+    mcp.run(transport="streamable-http")
+
+if __name__ == "__main__":
+    main()

ship_detector_mcp-0.1.0/src/ship_detector_mcp/__main__.py

@@ -0,0 +1,2 @@
+from ship_detector_mcp import main
+main()

ship_detector_mcp-0.1.0/src/ship_detector_mcp/server.py

@@ -0,0 +1,152 @@
+# server.py
+from mcp.server.fastmcp import FastMCP
+import os
+from ultralytics import YOLO
+import cv2
+import numpy as np
+import tempfile
+import base64
+from pathlib import Path
+from PIL import Image
+from io import BytesIO
+
+# 初始化 MCP 服务器，配置主机和端口
+mcp = FastMCP("ShipDetector", host="0.0.0.0", port=9000, stateless_http=True)
+
+# 加载 YOLO 模型（请替换为你的模型路径）
+
+import importlib.resources
+MODEL_PATH = os.getenv("MODEL_PATH", str(Path(__file__).parent / "weights" / "best.pt"))
+model = YOLO(MODEL_PATH)
+
+def load_image_from_input(image_input: str):
+    """
+    智能加载图片：自动判断输入是 Base64 还是文件路径
+
+    Args:
+        image_input: Base64 字符串 或 本地文件路径
+
+    Returns:
+        PIL Image 对象
+    """
+    # 判断是否是 Base64
+    is_base64 = False
+    # 检查是否是 PNG 的 Base64 头
+    if image_input.startswith('iVBORw0KGgo'):
+        is_base64 = True
+    # 检查是否是 JPG 的 Base64 头
+    elif image_input.startswith('/9j/4'):
+        is_base64 = True
+    # 检查是否是存在的文件路径
+    elif Path(image_input).exists():
+        pass  # 不是 Base64，是文件路径
+    else:
+        # 尝试作为 Base64 解码前100个字符
+        try:
+            base64.b64decode(image_input[:100])
+            is_base64 = True
+        except:
+            raise ValueError(f"输入既不是有效的 Base64，也不是存在的文件路径: {image_input[:50]}...")
+
+    if is_base64:
+        print("检测到 Base64 输入，正在解码...")
+        image_data = base64.b64decode(image_input)
+        image = Image.open(BytesIO(image_data)).convert("RGB")
+    else:
+        print(f"检测到文件路径输入: {image_input}")
+        image = Image.open(image_input).convert("RGB")
+
+    return image
+
+@mcp.tool()
+def detect_ship(
+        image_data: bytes = None,
+        image_path: str = None,
+        image_url: str = None,  # ← 新增
+        confidence: float = 0.5
+) -> str:
+    import requests
+    import tempfile
+    import os
+    import cv2
+    import numpy as np
+
+    temp_path = None
+
+    try:
+        # 情况1：有 image_data
+        if image_data is not None:
+            if isinstance(image_data, str):
+                import base64
+                image_data = base64.b64decode(image_data)
+
+            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+                tmp.write(image_data)
+                temp_path = tmp.name
+            image_to_use = temp_path
+            print(f"[INFO] 从 image_data 创建临时文件")
+
+        # 情况2：有 image_url（从 Dify 传来）
+        elif image_url is not None:
+            # 如果是相对路径，补全完整 URL
+            if image_url.startswith('/'):
+                image_url = f"http://localhost:5001{image_url}"
+
+            print(f"[INFO] 从 URL 下载图片: {image_url}")
+            response = requests.get(image_url, timeout=30)
+            response.raise_for_status()
+
+            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+                tmp.write(response.content)
+                temp_path = tmp.name
+            image_to_use = temp_path
+
+        # 情况3：有 image_path
+        elif image_path is not None:
+            if not os.path.exists(image_path):
+                return f"错误：图片文件不存在 - {image_path}"
+            image_to_use = image_path
+
+        else:
+            return "错误：请提供图片数据、URL或路径"
+
+        # 运行 YOLO 检测
+        results = model(image_to_use, conf=confidence)
+
+        # 处理结果
+        if not results or len(results) == 0:
+            return "未检测到船只"
+
+        detections = []
+        for box in results[0].boxes:
+            class_id = int(box.cls[0])
+            class_name = model.names[class_id]
+            conf = float(box.conf[0])
+            bbox = box.xyxy[0].tolist()
+            detections.append({
+                "class": class_name,
+                "confidence": round(conf, 4),
+                "bbox": [int(x) for x in bbox]
+            })
+
+        output = f"✅ 检测到 {len(detections)} 艘船:\n"
+        for i, ship in enumerate(detections, 1):
+            output += f"{i}. {ship['class']} - 置信度 {ship['confidence'] * 100:.2f}% - 位置 {ship['bbox']}\n"
+        return output
+
+    except Exception as e:
+        print(f"[ERROR] {e}")
+        return f"检测失败: {str(e)}"
+
+    finally:
+        if temp_path and os.path.exists(temp_path):
+            os.unlink(temp_path)
+@mcp.tool()
+def list_available_models() -> str:
+    """列出可用的检测模型"""
+    return f"当前模型: {MODEL_PATH}"
+
+
+if __name__ == "__main__":
+    # 使用 streamable-http 传输协议运行
+    mcp.run(transport="streamable-http")