MedSAM3

Use the infer_sam.py script to perform text-guided (concept-guided) medical image segmentation. You provide a medical concept as a text prompt, and the model segments the corresponding target.

Key Arguments:

• --config: Path to the configuration YAML file (e.g., configs/full_lora_config.yaml).
• --image: Path to the input medical image.
• --prompt: The medical concept/text description of the target (e.g., "skin lesion").
• --threshold: Confidence threshold for segmentation. Adjust this based on the modality (e.g., 0.8 for higher precision, 0.5 for higher sensitivity).
• --nms-iou: Non-Maximum Suppression IoU threshold.
• --output: Path where the resulting segmentation mask will be saved.

python3 infer_sam.py \
  --config configs/full_lora_config.yaml \
  --image path/to/image.jpg \
  --prompt "skin lesion" \
  --threshold 0.5 \
  --nms-iou 0.5 \
  --output skin_lesion.png

Use LoRAConfig to define how Low-Rank Adaptation (LoRA) is applied to a SAM3 model. You can specify the rank, scaling factor, dropout, and which specific module patterns should receive LoRA layers.

Parameters:

• rank (int): The rank of the LoRA matrices. Defaults to 4.
• alpha (float): The LoRA scaling factor. Defaults to 1.0.
• dropout (float): The dropout probability. Defaults to 0.0.
• target_modules (List[str]): A list of module name patterns to target. If None, it defaults to a comprehensive list covering standard attention, vision backbone (ViT-style), language backbone (CLIP-style), and transformer FFN layers. If 'all' is included in the list, it enables all possible modules including in_proj for MultiheadAttention.

Supported Target Module Patterns:

• Standard attention: q_proj, k_proj, v_proj, out_proj
• Vision backbone (ViT-style): qkv (fused), proj (output), fc1, fc2 (MLP)
• Language backbone (CLIP-style): c_fc, c_proj (MLP)
• Transformer FFN: linear1, linear2
• Special: all (enables everything listed above plus in_proj)

from src.lora.lora_utils import LoRAConfig

config = LoRAConfig(
    rank=8,
    alpha=2.0,
    dropout=0.1,
    target_modules=["q_proj", "v_proj"]
)

When initializing LoRATrainer, you can configure the following parameters to control the training process:

The LoRASAM3Dataset class is a PyTorch Dataset implementation designed for SAM3 LoRA fine-tuning using annotations in COCO format. It automatically handles loading images from a directory and mapping them to their corresponding annotations.

Key behaviors:

• Annotation Format: Expects a COCO-style JSON file.
• Filtering: Automatically skips images that have no associated annotations.
• Annotation Limiting: Uses max_ann_per_img to cap the number of annotations processed per image (default is 200) to prevent memory issues.
• Output: When indexed, it returns a dictionary containing the image (as a PIL Image or transformed tensor), annotations (list of dicts), and image_info (metadata).

from src.data.dataset import LoRASAM3Dataset

dataset = LoRASAM3Dataset(
    img_folder="path/to/images",
    ann_file="path/to/annotations.json",
    transforms=my_transforms_function,
    max_ann_per_img=100
)

sample = dataset[0]
# sample contains {'image': ..., 'annotations': ..., 'image_info': ...}

The create_dataloaders utility function simplifies the setup of the training and validation pipelines. It instantiates two LoRASAM3Dataset objects and wraps them in PyTorch DataLoader instances using a custom collate_fn to handle batching.

Arguments:

• train_img_folder / train_ann_file: Paths for training data.
• val_img_folder / val_ann_file: Paths for validation data.
• batch_size: Number of samples per batch (default: 1).
• num_workers: Number of subprocesses for data loading (default: 4).
• pin_memory: Whether to use pinned memory (default: True).
• train_transforms / val_transforms: Optional callable transforms for each set.
• max_ann_per_img: Maximum annotations per image (default: 200).

Returns:

• A tuple of (train_dataloader, val_dataloader).

from src.data.dataset import create_dataloaders

train_loader, val_loader = create_dataloaders(
    train_img_folder="data/train/imgs",
    train_ann_file="data/train/ann.json",
    val_img_folder="data/val/imgs",
    val_ann_file="data/val/ann.json",
    batch_size=4,
    train_transforms=my_train_transforms,
    val_transforms=my_val_transforms
)

These utilities allow you to manage the lifecycle of LoRA parameters independently of the base model weights.

• get_lora_state_dict(model): Returns a dictionary containing only the LoRA weights (lora_A and lora_B) for all LinearWithLoRA modules. This is useful for saving only the small adapter weights.
• load_lora_state_dict(model, state_dict): Loads LoRA weights from a provided dictionary into the model's LinearWithLoRA layers.
• merge_lora_weights(model): Permanently merges the LoRA weights into the base model's weights and converts LinearWithLoRA layers back into standard nn.Linear layers. This is typically used for inference deployment to avoid the overhead of LoRA computation.
• get_lora_parameters(model): Returns a list of all trainable LoRA parameters, useful for passing directly to an optimizer.

# To save only the adapters
state_dict = get_lora_state_dict(model)
torch.save(state_dict, "lora_weights.pt")

# To load adapters
state_dict = torch.load("lora_weights.pt")
load_lora_state_dict(model, state_dict)

# To merge for inference
model = merge_lora_weights(model)