Update MONAI/mcp_output/mcp_plugin/mcp_service.py

MONAI/mcp_output/mcp_plugin/mcp_service.py

@@ -1,133 +1,429 @@
+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="load_medical_dataset", description="Load a medical imaging dataset using MONAI")
-def load_medical_dataset(dataset_path: str) -> dict:
+
+@mcp.tool(name="get_monai_info", description="Get MONAI library information and configuration")
+def get_monai_info() -> dict:
     """
-    Load a medical imaging dataset using MONAI.
+    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.
 
     Parameters:
-    - dataset_path: Path to 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: Information about the loaded dataset.
+    - dict: Model information and parameter count.
     """
     try:
-        from monai.data import Dataset
-        from monai.transforms import LoadImaged
+        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
+        )
         
-        data = [{"image": dataset_path}]
-        transforms = LoadImaged(keys=["image"])
-        dataset = Dataset(data, transforms)
+        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,
-            "dataset": str(dataset)
+            "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="train_segmentation_model", description="Train a segmentation model using MONAI")
-def train_segmentation_model(config: dict) -> 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:
     """
-    Train a segmentation model using MONAI.
+    Create a SegResNet model for medical image segmentation.
 
     Parameters:
-    - config: Configuration dictionary for training.
+    - 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: Training results.
+    - dict: Model information and parameter count.
     """
     try:
-        from monai.engines import SupervisedTrainer
-        from monai.transforms import Compose
+        model = SegResNet(
+            spatial_dims=spatial_dims,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            init_filters=init_filters
+        )
         
-        # Example: Initialize trainer with config
-        trainer = SupervisedTrainer(**config)
-        trainer.run()
+        total_params = sum(p.numel() for p in model.parameters())
         
         return {
             "success": True,
-            "message": "Training completed successfully."
+            "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="evaluate_model", description="Evaluate a trained model using MONAI")
-def evaluate_model(model_path: str, test_data: list) -> 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:
     """
-    Evaluate a trained model using MONAI.
+    Compute Dice similarity coefficient between prediction and ground truth.
 
     Parameters:
-    - model_path: Path to the trained model.
-    - test_data: Test dataset.
+    - 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: Evaluation metrics.
+    - dict: Dice score for each class.
     """
     try:
-        from monai.engines import SupervisedEvaluator
+        import torch
         
-        evaluator = SupervisedEvaluator(model_path=model_path, data=test_data)
-        metrics = evaluator.run()
+        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,
-            "metrics": metrics
+            "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="apply_transforms", description="Apply MONAI transforms to medical images")
-def apply_transforms(image_path: str, transforms: list) -> 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:
     """
-    Apply MONAI transforms to medical images.
+    Compute segmentation loss between prediction and ground truth.
 
     Parameters:
-    - image_path: Path to the image.
-    - transforms: List of transforms to apply.
+    - 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: Transformed image data.
+    - dict: Loss value.
     """
     try:
-        from monai.transforms import Compose
+        import torch
+        
+        pred_tensor = torch.tensor(prediction).unsqueeze(0).unsqueeze(0).float()
+        gt_tensor = torch.tensor(ground_truth).unsqueeze(0).unsqueeze(0).float()
         
-        composed_transforms = Compose(transforms)
-        transformed_image = composed_transforms(image_path)
+        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,
-            "transformed_image": transformed_image
+            "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="visualize_segmentation", description="Visualize segmentation results using MONAI")
-def visualize_segmentation(image_path: str, segmentation_path: str) -> 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:
     """
-    Visualize segmentation results using MONAI.
+    Apply intensity transforms to a medical image.
 
     Parameters:
-    - image_path: Path to the original image.
-    - segmentation_path: Path to the segmentation result.
+    - 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: Visualization status.
+    - dict: Transformed image statistics.
     """
     try:
-        from monai.visualize import blend_images
+        img_array = np.array(image, dtype=np.float32)
         
-        blended_image = blend_images(image_path, segmentation_path)
+        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,
-            "blended_image": blended_image
+            "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="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)}
+
 
 def create_app() -> FastMCP:
     """