Update MONAI/mcp_output/mcp_plugin/mcp_service.py

MONAI/mcp_output/mcp_plugin/mcp_service.py

@@ -1,5 +1,6 @@
 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")
@@ -7,127 +8,422 @@ if source_path not in sys.path:
     sys.path.insert(0, source_path)
 
 from fastmcp import FastMCP
+import numpy as np
 
-# Import core modules from the local source directory
-from monai.data import Dataset, DataLoader
-from monai.transforms import Compose, LoadImage, NormalizeIntensity
-from monai.networks.nets import UNet
-from monai.engines import SupervisedTrainer
-from monai.metrics import DiceMetric
+# 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="load_dataset", description="Load a dataset using MONAI's Dataset class")
-def load_dataset(data_dir: str) -> dict:
+
+@mcp.tool(name="get_monai_info", description="Get MONAI library information and configuration")
+def get_monai_info() -> 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:
     """
-    Load a dataset from the specified directory.
+    Create a UNet model for medical image segmentation.
 
     Parameters:
-    - data_dir (str): The directory containing the dataset.
+    - 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]).
 
     Returns:
-    - dict: A dictionary containing success status and the dataset object or error message.
+    - dict: Model information and parameter count.
     """
     try:
-        dataset = Dataset(data_dir)
-        return {"success": True, "result": dataset}
+        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
+        )
+        
+        total_params = sum(p.numel() for p in model.parameters())
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        
+        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
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="create_dataloader", description="Create a DataLoader for the dataset")
-def create_dataloader(dataset: Dataset, batch_size: int) -> dict:
+
+@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:
     """
-    Create a DataLoader for the given dataset.
+    Create a SegResNet model for medical image segmentation.
 
     Parameters:
-    - dataset (Dataset): The dataset to load.
-    - batch_size (int): The number of samples per batch.
+    - 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.
 
     Returns:
-    - dict: A dictionary containing success status and the DataLoader object or error message.
+    - dict: Model information and parameter count.
     """
     try:
-        dataloader = DataLoader(dataset, batch_size=batch_size)
-        return {"success": True, "result": dataloader}
+        model = SegResNet(
+            spatial_dims=spatial_dims,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            init_filters=init_filters
+        )
+        
+        total_params = sum(p.numel() for p in model.parameters())
+        
+        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
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="apply_transforms", description="Apply transforms to the dataset")
-def apply_transforms(dataset: Dataset) -> dict:
+@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:
     """
-    Apply a series of transforms to the dataset.
+    Compute Dice similarity coefficient between prediction and ground truth.
 
     Parameters:
-    - dataset (Dataset): The dataset to transform.
+    - prediction: Predicted segmentation as nested list (2D or 3D).
+    - ground_truth: Ground truth segmentation as nested list.
+    - include_background: Whether to include background class.
 
     Returns:
-    - dict: A dictionary containing success status and the transformed dataset or error message.
+    - dict: Dice score for each class.
     """
     try:
-        transforms = Compose([LoadImage(), NormalizeIntensity()])
-        transformed_dataset = [transforms(item) for item in dataset]
-        return {"success": True, "result": transformed_dataset}
+        import torch
+        
+        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()
+        
+        return {
+            "success": True,
+            "result": {
+                "dice_score": dice_score,
+                "include_background": include_background
+            },
+            "error": None
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="initialize_unet", description="Initialize a UNet model")
-def initialize_unet(spatial_dims: int, in_channels: int, out_channels: int) -> dict:
+
+@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:
     """
-    Initialize a UNet model with the specified parameters.
+    Compute segmentation loss between prediction and ground truth.
 
     Parameters:
-    - spatial_dims (int): The number of spatial dimensions.
-    - in_channels (int): The number of input channels.
-    - out_channels (int): The number of output channels.
+    - prediction: Predicted logits as nested list.
+    - ground_truth: Ground truth labels as nested list.
+    - loss_type: Type of loss ('dice', 'dice_ce', 'focal', 'tversky').
 
     Returns:
-    - dict: A dictionary containing success status and the UNet model or error message.
+    - dict: Loss value.
     """
     try:
-        model = UNet(spatial_dims=spatial_dims, in_channels=in_channels, out_channels=out_channels)
-        return {"success": True, "result": model}
+        import torch
+        
+        pred_tensor = torch.tensor(prediction).unsqueeze(0).unsqueeze(0).float()
+        gt_tensor = torch.tensor(ground_truth).unsqueeze(0).unsqueeze(0).float()
+        
+        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()
+        
+        return {
+            "success": True,
+            "result": {
+                "loss_type": loss_type,
+                "loss_value": loss_value
+            },
+            "error": None
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="train_model", description="Train a model using MONAI's SupervisedTrainer")
-def train_model(model, dataloader: DataLoader, max_epochs: int) -> dict:
+@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:
     """
-    Train a model using the specified dataloader and number of epochs.
+    Apply intensity transforms to a medical image.
 
     Parameters:
-    - model: The model to train.
-    - dataloader (DataLoader): The DataLoader for training data.
-    - max_epochs (int): The maximum number of training epochs.
+    - 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.
 
     Returns:
-    - dict: A dictionary containing success status and training results or error message.
+    - dict: Transformed image statistics.
     """
     try:
-        trainer = SupervisedTrainer(max_epochs=max_epochs, train_data_loader=dataloader, network=model)
-        trainer.run()
-        return {"success": True, "result": "Training completed"}
+        img_array = np.array(image, dtype=np.float32)
+        
+        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
+        
+        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
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="calculate_dice", description="Calculate Dice metric for model evaluation")
-def calculate_dice(predictions, targets) -> dict:
+
+@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:
     """
-    Calculate the Dice metric for the given predictions and targets.
+    Resize a medical image to target spatial size.
 
     Parameters:
-    - predictions: The model predictions.
-    - targets: The ground truth targets.
+    - 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: A dictionary containing success status and the Dice score or error message.
+    - dict: Resized image info.
     """
     try:
-        dice_metric = DiceMetric()
-        dice_score = dice_metric(predictions, targets)
-        return {"success": True, "result": dice_score}
+        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, "error": str(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)}
+
 
 def create_app() -> FastMCP:
     """