Update MONAI/mcp_output/mcp_plugin/mcp_service.py

MONAI/mcp_output/mcp_plugin/mcp_service.py

@@ -1,429 +1,133 @@
-import os
-import sys
-from typing import List, Optional, Dict, Any
-
-# Add the local source directory to sys.path
-source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
-if source_path not in sys.path:
-    sys.path.insert(0, source_path)
-
 from fastmcp import FastMCP
-import numpy as np
-
-# Import MONAI modules
-import monai
-from monai.transforms import (
-    Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, 
-    NormalizeIntensity, Resize, RandRotate, RandFlip, ToTensor
-)
-from monai.networks.nets import UNet, VNet, AttentionUnet, SegResNet, DenseNet121
-from monai.losses import DiceLoss, DiceCELoss, FocalLoss, TverskyLoss
-from monai.metrics import DiceMetric, MeanIoU, HausdorffDistanceMetric
-from monai.data import decollate_batch
-from monai.inferers import sliding_window_inference
 
 # Create the FastMCP service application
 mcp = FastMCP("monai_service")
 
-
-@mcp.tool(name="get_monai_info", description="Get MONAI library information and configuration")
-def get_monai_info() -> dict:
+@mcp.tool(name="load_medical_dataset", description="Load a medical imaging dataset using MONAI")
+def load_medical_dataset(dataset_path: str) -> dict:
     """
-    Get MONAI library version and system configuration.
-
-    Returns:
-    - dict: MONAI version and configuration info.
-    """
-    try:
-        config = monai.config.get_config_values()
-        return {
-            "success": True, 
-            "result": {
-                "version": monai.__version__,
-                "config": {k: str(v) for k, v in config.items()}
-            },
-            "error": None
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="create_unet_model", description="Create and initialize a UNet model")
-def create_unet_model(
-    spatial_dims: int = 3,
-    in_channels: int = 1,
-    out_channels: int = 2,
-    channels: List[int] = None,
-    strides: List[int] = None
-) -> dict:
-    """
-    Create a UNet model for medical image segmentation.
+    Load a medical imaging dataset using MONAI.
 
     Parameters:
-    - spatial_dims: Number of spatial dimensions (2 or 3).
-    - in_channels: Number of input channels.
-    - out_channels: Number of output channels/classes.
-    - channels: Feature channels per layer (default: [16, 32, 64, 128, 256]).
-    - strides: Strides per layer (default: [2, 2, 2, 2]).
+    - dataset_path: Path to the dataset.
 
     Returns:
-    - dict: Model information and parameter count.
+    - dict: Information about the loaded dataset.
     """
     try:
-        if channels is None:
-            channels = [16, 32, 64, 128, 256]
-        if strides is None:
-            strides = [2, 2, 2, 2]
-        
-        model = UNet(
-            spatial_dims=spatial_dims,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            channels=channels,
-            strides=strides,
-            num_res_units=2
-        )
+        from monai.data import Dataset
+        from monai.transforms import LoadImaged
         
-        total_params = sum(p.numel() for p in model.parameters())
-        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        data = [{"image": dataset_path}]
+        transforms = LoadImaged(keys=["image"])
+        dataset = Dataset(data, transforms)
         
         return {
             "success": True,
-            "result": {
-                "model_type": "UNet",
-                "spatial_dims": spatial_dims,
-                "in_channels": in_channels,
-                "out_channels": out_channels,
-                "channels": channels,
-                "strides": strides,
-                "total_parameters": total_params,
-                "trainable_parameters": trainable_params
-            },
-            "error": None
+            "dataset": str(dataset)
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
-
-@mcp.tool(name="create_segresnet_model", description="Create a SegResNet model")
-def create_segresnet_model(
-    spatial_dims: int = 3,
-    in_channels: int = 1,
-    out_channels: int = 2,
-    init_filters: int = 16
-) -> dict:
+@mcp.tool(name="train_segmentation_model", description="Train a segmentation model using MONAI")
+def train_segmentation_model(config: dict) -> dict:
     """
-    Create a SegResNet model for medical image segmentation.
+    Train a segmentation model using MONAI.
 
     Parameters:
-    - spatial_dims: Number of spatial dimensions (2 or 3).
-    - in_channels: Number of input channels.
-    - out_channels: Number of output channels/classes.
-    - init_filters: Initial number of filters.
+    - config: Configuration dictionary for training.
 
     Returns:
-    - dict: Model information and parameter count.
+    - dict: Training results.
     """
     try:
-        model = SegResNet(
-            spatial_dims=spatial_dims,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            init_filters=init_filters
-        )
+        from monai.engines import SupervisedTrainer
+        from monai.transforms import Compose
         
-        total_params = sum(p.numel() for p in model.parameters())
+        # Example: Initialize trainer with config
+        trainer = SupervisedTrainer(**config)
+        trainer.run()
         
         return {
             "success": True,
-            "result": {
-                "model_type": "SegResNet",
-                "spatial_dims": spatial_dims,
-                "in_channels": in_channels,
-                "out_channels": out_channels,
-                "init_filters": init_filters,
-                "total_parameters": total_params
-            },
-            "error": None
+            "message": "Training completed successfully."
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="compute_dice_score", description="Compute Dice score between prediction and ground truth")
-def compute_dice_score(
-    prediction: List[List[List[float]]],
-    ground_truth: List[List[List[float]]],
-    include_background: bool = False
-) -> dict:
+@mcp.tool(name="evaluate_model", description="Evaluate a trained model using MONAI")
+def evaluate_model(model_path: str, test_data: list) -> dict:
     """
-    Compute Dice similarity coefficient between prediction and ground truth.
+    Evaluate a trained model using MONAI.
 
     Parameters:
-    - prediction: Predicted segmentation as nested list (2D or 3D).
-    - ground_truth: Ground truth segmentation as nested list.
-    - include_background: Whether to include background class.
+    - model_path: Path to the trained model.
+    - test_data: Test dataset.
 
     Returns:
-    - dict: Dice score for each class.
+    - dict: Evaluation metrics.
     """
     try:
-        import torch
+        from monai.engines import SupervisedEvaluator
         
-        pred_tensor = torch.tensor(prediction).unsqueeze(0).unsqueeze(0).float()
-        gt_tensor = torch.tensor(ground_truth).unsqueeze(0).unsqueeze(0).float()
-        
-        dice_metric = DiceMetric(include_background=include_background, reduction="mean")
-        dice_metric(y_pred=pred_tensor, y=gt_tensor)
-        dice_score = dice_metric.aggregate().item()
-        dice_metric.reset()
+        evaluator = SupervisedEvaluator(model_path=model_path, data=test_data)
+        metrics = evaluator.run()
         
         return {
             "success": True,
-            "result": {
-                "dice_score": dice_score,
-                "include_background": include_background
-            },
-            "error": None
+            "metrics": metrics
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="compute_loss", description="Compute segmentation loss")
-def compute_loss(
-    prediction: List[List[List[float]]],
-    ground_truth: List[List[List[float]]],
-    loss_type: str = "dice"
-) -> dict:
+@mcp.tool(name="apply_transforms", description="Apply MONAI transforms to medical images")
+def apply_transforms(image_path: str, transforms: list) -> dict:
     """
-    Compute segmentation loss between prediction and ground truth.
+    Apply MONAI transforms to medical images.
 
     Parameters:
-    - prediction: Predicted logits as nested list.
-    - ground_truth: Ground truth labels as nested list.
-    - loss_type: Type of loss ('dice', 'dice_ce', 'focal', 'tversky').
+    - image_path: Path to the image.
+    - transforms: List of transforms to apply.
 
     Returns:
-    - dict: Loss value.
+    - dict: Transformed image data.
     """
     try:
-        import torch
-        
-        pred_tensor = torch.tensor(prediction).unsqueeze(0).unsqueeze(0).float()
-        gt_tensor = torch.tensor(ground_truth).unsqueeze(0).unsqueeze(0).float()
+        from monai.transforms import Compose
         
-        if loss_type == "dice":
-            loss_fn = DiceLoss(sigmoid=True)
-        elif loss_type == "dice_ce":
-            loss_fn = DiceCELoss(sigmoid=True)
-        elif loss_type == "focal":
-            loss_fn = FocalLoss()
-        elif loss_type == "tversky":
-            loss_fn = TverskyLoss(sigmoid=True)
-        else:
-            return {"success": False, "result": None, "error": f"Unknown loss type: {loss_type}"}
-        
-        loss_value = loss_fn(pred_tensor, gt_tensor).item()
+        composed_transforms = Compose(transforms)
+        transformed_image = composed_transforms(image_path)
         
         return {
             "success": True,
-            "result": {
-                "loss_type": loss_type,
-                "loss_value": loss_value
-            },
-            "error": None
+            "transformed_image": transformed_image
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="apply_intensity_transforms", description="Apply intensity transforms to an image")
-def apply_intensity_transforms(
-    image: List[List[List[float]]],
-    normalize: bool = True,
-    scale_intensity: bool = False,
-    target_min: float = 0.0,
-    target_max: float = 1.0
-) -> dict:
+@mcp.tool(name="visualize_segmentation", description="Visualize segmentation results using MONAI")
+def visualize_segmentation(image_path: str, segmentation_path: str) -> dict:
     """
-    Apply intensity transforms to a medical image.
+    Visualize segmentation results using MONAI.
 
     Parameters:
-    - image: Input image as nested list (2D or 3D).
-    - normalize: Whether to normalize intensity (zero mean, unit std).
-    - scale_intensity: Whether to scale intensity to target range.
-    - target_min: Minimum value for scaling.
-    - target_max: Maximum value for scaling.
+    - image_path: Path to the original image.
+    - segmentation_path: Path to the segmentation result.
 
     Returns:
-    - dict: Transformed image statistics.
+    - dict: Visualization status.
     """
     try:
-        img_array = np.array(image, dtype=np.float32)
+        from monai.visualize import blend_images
         
-        transforms_list = []
-        applied_transforms = []
-        
-        if normalize:
-            transforms_list.append(NormalizeIntensity())
-            applied_transforms.append("NormalizeIntensity")
-        
-        if scale_intensity:
-            transforms_list.append(ScaleIntensity(minv=target_min, maxv=target_max))
-            applied_transforms.append(f"ScaleIntensity({target_min}, {target_max})")
-        
-        if transforms_list:
-            transform = Compose(transforms_list)
-            result = transform(img_array)
-        else:
-            result = img_array
+        blended_image = blend_images(image_path, segmentation_path)
         
         return {
             "success": True,
-            "result": {
-                "applied_transforms": applied_transforms,
-                "output_shape": list(result.shape),
-                "output_min": float(result.min()),
-                "output_max": float(result.max()),
-                "output_mean": float(result.mean()),
-                "output_std": float(result.std())
-            },
-            "error": None
+            "blended_image": blended_image
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="resize_image", description="Resize a medical image using MONAI")
-def resize_image(
-    image: List[List[List[float]]],
-    spatial_size: List[int],
-    mode: str = "trilinear"
-) -> dict:
-    """
-    Resize a medical image to target spatial size.
-
-    Parameters:
-    - image: Input image as nested list (2D or 3D).
-    - spatial_size: Target spatial size [H, W] or [D, H, W].
-    - mode: Interpolation mode ('nearest', 'bilinear', 'trilinear').
-
-    Returns:
-    - dict: Resized image info.
-    """
-    try:
-        img_array = np.array(image, dtype=np.float32)
-        original_shape = img_array.shape
-        
-        # Add channel dimension if needed
-        if len(img_array.shape) == len(spatial_size):
-            img_array = img_array[np.newaxis, ...]
-        
-        resize_transform = Resize(spatial_size=spatial_size, mode=mode)
-        resized = resize_transform(img_array)
-        
-        return {
-            "success": True,
-            "result": {
-                "original_shape": list(original_shape),
-                "target_size": spatial_size,
-                "output_shape": list(resized.shape),
-                "interpolation_mode": mode
-            },
-            "error": None
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="get_network_architectures", description="Get details about available network architectures")
-def get_network_architectures() -> dict:
-    """
-    Get detailed information about available MONAI network architectures.
-
-    Returns:
-    - dict: Network architecture details.
-    """
-    try:
-        architectures = {
-            "UNet": {
-                "description": "U-Net architecture for semantic segmentation",
-                "use_case": "General medical image segmentation",
-                "parameters": ["spatial_dims", "in_channels", "out_channels", "channels", "strides"]
-            },
-            "VNet": {
-                "description": "V-Net for volumetric medical image segmentation",
-                "use_case": "3D medical image segmentation",
-                "parameters": ["spatial_dims", "in_channels", "out_channels"]
-            },
-            "AttentionUnet": {
-                "description": "Attention U-Net with attention gates",
-                "use_case": "Segmentation with attention mechanism",
-                "parameters": ["spatial_dims", "in_channels", "out_channels", "channels", "strides"]
-            },
-            "SegResNet": {
-                "description": "ResNet-based encoder-decoder for segmentation",
-                "use_case": "Medical image segmentation with residual connections",
-                "parameters": ["spatial_dims", "in_channels", "out_channels", "init_filters"]
-            },
-            "SwinUNETR": {
-                "description": "Swin Transformer based U-Net",
-                "use_case": "State-of-the-art medical image segmentation",
-                "parameters": ["img_size", "in_channels", "out_channels", "feature_size"]
-            },
-            "DenseNet121": {
-                "description": "DenseNet for classification",
-                "use_case": "Medical image classification",
-                "parameters": ["spatial_dims", "in_channels", "out_channels"]
-            }
-        }
-        
-        return {"success": True, "result": architectures, "error": None}
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="create_densenet_classifier", description="Create a DenseNet classifier model")
-def create_densenet_classifier(
-    spatial_dims: int = 3,
-    in_channels: int = 1,
-    out_channels: int = 2
-) -> dict:
-    """
-    Create a DenseNet121 classifier for medical image classification.
-
-    Parameters:
-    - spatial_dims: Number of spatial dimensions (2 or 3).
-    - in_channels: Number of input channels.
-    - out_channels: Number of output classes.
-
-    Returns:
-    - dict: Model information.
-    """
-    try:
-        model = DenseNet121(
-            spatial_dims=spatial_dims,
-            in_channels=in_channels,
-            out_channels=out_channels
-        )
-        
-        total_params = sum(p.numel() for p in model.parameters())
-        
-        return {
-            "success": True,
-            "result": {
-                "model_type": "DenseNet121",
-                "spatial_dims": spatial_dims,
-                "in_channels": in_channels,
-                "out_channels": out_channels,
-                "total_parameters": total_params
-            },
-            "error": None
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
 def create_app() -> FastMCP:
     """