NeuroSeg-MF: robust neuron segmentation in two-photon Ca2+ imaging using multi-feature fusion and detection-guided SAM.

Code ↔ Paper

11 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 11 matches

[1] § Materials and methods › Framework of NeuroSeg-MF › Multi-feature detection network ↔ ultralytics/nn/tasks.py, lines 1778–1920 · score 0.80 · basic block, GateFusion, A2C2f, RT DETR, AIFI, global
[2] § Materials and methods › Framework of NeuroSeg-MF › Multi-feature detection network ↔ ultralytics/nn/modules/__init__.py, lines 37–71 · score 0.80 · RepConv, basic block, GateFusion, A2C2f, AIFI, decoder
[3] § Materials and methods › Image preprocessing › Data augmentation ↔ ultralytics/data/augment.py, lines 2649–2750 · score 0.76 · horizontal flipping, Random erasing, jitter, augmentation, HSV, probability
[4] § Materials and methods › Framework of NeuroSeg-MF › Multi-feature detection network ↔ ultralytics/nn/modules/__init__.py, lines 37–71 · score 0.74 · RepConv, BasicBlock, GateFusion, A2C2f, AIFI, decoder
[5] § Materials and methods › Framework of NeuroSeg-MF › Detection-guided SAM for neuron segmentation ↔ tools/SAM_box_to_mask.py, lines 1–65 · score 0.74 · candidate masks generated, binary masks, box prompt, segmentation masks, overlays, bounding boxes
[6] § Materials and methods › Framework of NeuroSeg-MF › Multi-feature detection network ↔ ultralytics/nn/tasks.py, lines 1778–1920 · score 0.74 · BasicBlock, GateFusion, A2C2f, RT DETR, AIFI, architecture
[7] § Materials and methods › Framework of NeuroSeg-MF › Detection-guided SAM for neuron segmentation ↔ tools/SAM_box_to_mask.py, lines 1–65 · score 0.64 · generate candidate masks, box prompt, detection box, scores, predicted, neuron
[8] § Materials and methods › Image preprocessing › Pseudo-depth map ↔ tools/generate_depth_twophoton.py, lines 61–145 · score 0.56 · bit grayscale, pseudo depth map
[9] § Materials and methods › Evaluation metrics ↔ ultralytics/utils/coco_metrics.py, lines 282–404 · score 0.51 · IoU threshold, F1 score, recall, precision, metrics, predicted
[10] § Materials and methods › Image preprocessing › Pseudo-depth map ↔ tools/generate_depth_twophoton.py, lines 61–145 · score 0.50 · pseudo depth map, smoothing, V2, filtering, clipping, resized
[11] § Materials and methods › Image preprocessing › Correlation map ↔ tools/generate_corr.py, lines 404–483 · score 0.50 · local neighborhood, correlation map, frames, pixels

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The paper is loaded by your browser, with JavaScript: read it at doi.org or on Europe PMC.

The authors' code

Python · 2,035 lines · 82 KB · no license · 2 matches

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import contextlib
import ast
import pickle
import re
import types
from copy import deepcopy
from pathlib import Path
import torch
import torch.nn as nn
# Extra modules conditional import(defaulttext)
# from ultralytics.nn.extra_modules import *
# from ultralytics.nn.backbone.convnextv2 import *
# from ultralytics.nn.backbone.fasternet import *
# from ultralytics.nn.backbone.efficientViT import *
# from ultralytics.nn.backbone.EfficientFormerV2 import *
# from ultralytics.nn.backbone.VanillaNet import *
# from ultralytics.nn.backbone.revcol import *
# from ultralytics.nn.backbone.lsknet import *
# from ultralytics.nn.backbone.SwinTransformer import *
# from ultralytics.nn.backbone.repvit import *
# from ultralytics.nn.backbone.CSwomTramsformer import *
# from ultralytics.nn.backbone.UniRepLKNet import *
# from ultralytics.nn.backbone.TransNext import *
# from ultralytics.nn.backbone.rmt import *
# from ultralytics.nn.backbone.pkinet import *
# from ultralytics.nn.backbone.mobilenetv4 import *
# from ultralytics.nn.backbone.starnet import *
# from ultralytics.nn.backbone.inceptionnext import *
# from ultralytics.nn.extra_modules.mobileMamba.mobilemamba import *
# from ultralytics.nn.backbone.MambaOut import *
# from ultralytics.nn.backbone.overlock import *
# from ultralytics.nn.backbone.lsnet import *
# except:
# pass
from ultralytics.nn.autobackend import check_class_names
from ultralytics.nn.modules import (
AIFI,
A2C2f,
BasicBlock,
Blocks,
Bottleneck,
C3,
Concat,
Conv,
ConvNormLayer,
DWConv,
GateFusion,
HGBlock,
HGStem,
Index,
RepC3,
RepConv,
RTDETRBottleNeck,
RTDETRDecoder,
get_activation,
)
from ultralytics.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, YAML, colorstr, emojis
from ultralytics.utils.checks import check_requirements, check_suffix, check_yaml
from ultralytics.utils.loss import v8DetectionLoss
from ultralytics.utils.ops import make_divisible
from ultralytics.utils.patches import torch_load
from ultralytics.utils.plotting import feature_visualization
from ultralytics.utils.torch_utils import (
fuse_conv_and_bn,
initialize_weights,
intersect_dicts,
model_info,
scale_img,
smart_inference_mode,
time_sync,
)
try:
from ultralytics.nn.mm import MultiModalRouter, MultiModalConfigParser
HookManager = None
MULTIMODAL_AVAILABLE = True
except Exception:
MULTIMODAL_AVAILABLE = False
MultiModalRouter = MultiModalConfigParser = None
HookManager = None
CONTRAST_AVAILABLE = False
DETECT_CLASS: tuple = ()
SEGMENT_CLASS: tuple = ()
POSE_CLASS: tuple = ()
OBB_CLASS: tuple = ()
C3K2_CLASS: tuple = ()
C2PSA_CLASS: tuple = ()
SPPF_CLASS: tuple = ()
NECK_CLASS: tuple = ()
LSCD_AVAILABLE = False
C3K2_EXTRACTION_AVAILABLE = False
SPPF_EXTRACTION_AVAILABLE = False
C2PSA_EXTRACTION_AVAILABLE = False
NECK_EXTRACTION_AVAILABLE = False
class _UnsupportedModule(torch.nn.Module):
"""Placeholder for modules removed from the NeuroSeg-MF minimal runtime."""
def __init__(self, *args, **kwargs):
super().__init__()
raise NotImplementedError("This module is not included in the NeuroSeg-MF minimal runtime.")
# Compatibility aliases referenced by unused upstream methods. They are not used by NeuroSeg-MF.
Conv2 = ConvTranspose = DWConvTranspose2d = GhostConv = GhostBottleneck = Focus = BottleneckCSP = _UnsupportedModule
C1 = C2 = C2f = C2fAttn = C2fCIB = C2PSA = C3TR = C3Ghost = C3k2 = C3x = _UnsupportedModule
AConv = ADown = ELAN1 = PSA = SPP = SPPF = SPPELAN = SCDown = _UnsupportedModule
RepNCSPELAN4 = RepVGGDW = ResNetLayer = TorchVision = CBLinear = CBFuse = _UnsupportedModule
Classify = Detect = v8Detect = Segment = Pose = OBB = WorldDetect = YOLOEDetect = YOLOESegment = v10Detect = ImagePoolingAttn = LRPCHead = _UnsupportedModule
MutilScaleEdgeInfoGenetator = ConvEdgeFusion = GetIndexOutput = MCFGatedFusion = _UnsupportedModule
class BaseModel(torch.nn.Module):
"""
Base class for all YOLO models in the Ultralytics family.
This class provides common functionality for YOLO models including forward pass handling, model fusion,
information display, and weight loading capabilities.
Attributes:
model (torch.nn.Module): The neural network model.
save (list): List of layer indices to save outputs from.
stride (torch.Tensor): Model stride values.
Methods:
forward: Perform forward pass for training or inference.
predict: Perform inference on input tensor.
fuse: Fuse Conv2d and BatchNorm2d layers for optimization.
info: Print model information.
load: Load weights into the model.
loss: Compute loss for training.
Examples:
Create a BaseModel instance
>>> model = BaseModel()
>>> model.info() # Display model information
"""
def forward(self, x, *args, **kwargs):
"""
Perform forward pass of the model for either training or inference.
If x is a dict, calculates and returns the loss for training. Otherwise, returns predictions for inference.
Args:
x (torch.Tensor | dict): Input tensor for inference, or dict with image tensor and labels for training.
*args (Any): Variable length argument list.
**kwargs (Any): Arbitrary keyword arguments.
Returns:
(torch.Tensor): Loss if x is a dict (training), or network predictions (inference).
"""
if isinstance(x, dict): # for cases of training and validating while training.
return self.loss(x, *args, **kwargs)
return self.predict(x, *args, **kwargs)
def predict(self, x, profile=False, visualize=False, augment=False, embed=None):
"""
Perform a forward pass through the network.
Args:
x (torch.Tensor): The input tensor to the model.
profile (bool): Print the computation time of each layer if True.
visualize (bool): Save the feature maps of the model if True.
augment (bool): Augment image during prediction.
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(torch.Tensor): The last output of the model.
"""
if augment:
return self._predict_augment(x)
return self._predict_once(x, profile, visualize, embed)
def _predict_once(self, x, profile=False, visualize=False, embed=None):
"""
Perform a forward pass through the network.
Args:
x (torch.Tensor): The input tensor to the model.
profile (bool): Print the computation time of each layer if True.
visualize (bool): Save the feature maps of the model if True.
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(torch.Tensor): The last output of the model.
"""
# ===== MULTIMODAL EXTENSION START - multi-modaltext =====
mm_router = None
mm_routing_enabled = False
mm_input_sources = None
# Check if this model has a persistent router (RTDETRDetectionModel)
if hasattr(self, 'mm_router') and self.mm_router is not None:
# Use persistent router from model initialization
mm_router = self.mm_router
mm_routing_enabled, mm_input_sources = mm_router.setup_multimodal_routing(x, profile)
if profile:
LOGGER.info("MultiModal: textrouter")
elif MULTIMODAL_AVAILABLE:
try:
from ultralytics.nn.mm import MultiModalRouter
# Create temporary router for other model types
config_dict = getattr(self, 'yaml', None)
mm_router = MultiModalRouter(config_dict, verbose=profile)
mm_routing_enabled, mm_input_sources = mm_router.setup_multimodal_routing(x, profile)
if profile:
LOGGER.info("MultiModal: textrouter")
except Exception as e:
if profile:
LOGGER.warning(f"MultiModal routing initialization failed: {e}")
# ===== MULTIMODAL EXTENSION END =====
y, dt, embeddings = [], [], [] # outputs
embed = frozenset(embed) if embed is not None else {-1}
max_idx = max(embed)
for m in self.model:
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
# ===== MULTIMODAL EXTENSION START - multi-modaltext =====
# Apply multimodal routing if enabled and module has MM attributes
if mm_routing_enabled and mm_input_sources and mm_router:
routed_x = mm_router.route_layer_input(x, m, mm_input_sources, profile)
if routed_x is not None:
x = routed_x
# Check for spatial reset requirement
if mm_router and hasattr(m, '_mm_spatial_reset') and m._mm_spatial_reset:
x = mm_router.reset_spatial_input(x, m, mm_input_sources, profile)
# ===== MULTIMODAL EXTENSION END =====
if profile:
self._profile_one_layer(m, x, dt)
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
if m.i in embed:
embeddings.append(torch.nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1)) # flatten
if m.i == max_idx:
return torch.unbind(torch.cat(embeddings, 1), dim=0)
return x
def _predict_augment(self, x):
"""Perform augmentations on input image x and return augmented inference."""
LOGGER.warning(
f"{self.__class__.__name__} does not support 'augment=True' prediction. "
f"Reverting to single-scale prediction."
)
return self._predict_once(x)
def _profile_one_layer(self, m, x, dt):
"""
Profile the computation time and FLOPs of a single layer of the model on a given input.
Args:
m (torch.nn.Module): The layer to be profiled.
x (torch.Tensor): The input data to the layer.
dt (list): A list to store the computation time of the layer.
"""
try:
import thop
except ImportError:
thop = None # conda support without 'ultralytics-thop' installed
c = m == self.model[-1] and isinstance(x, list) # is final layer list, copy input as inplace fix
flops = thop.profile(m, inputs=[x.copy() if c else x], verbose=False)[0] / 1e9 * 2 if thop else 0 # GFLOPs
t = time_sync()
for _ in range(10):
m(x.copy() if c else x)
dt.append((time_sync() - t) * 100)
if m == self.model[0]:
LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s} module")
LOGGER.info(f"{dt[-1]:10.2f} {flops:10.2f} {m.np:10.0f} {m.type}")
if c:
LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total")
def fuse(self, verbose=True):
"""
Fuse the `Conv2d()` and `BatchNorm2d()` layers of the model into a single layer for improved computation
efficiency.
Returns:
(torch.nn.Module): The fused model is returned.
"""
if not self.is_fused():
for m in self.model.modules():
if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"):
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
delattr(m, "bn") # remove batchnorm
m.forward = m.forward_fuse # update forward
if isinstance(m, RepConv):
m.fuse_convs()
m.forward = m.forward_fuse # update forward
self.info(verbose=verbose)
return self
def is_fused(self, thresh=10):
"""
Check if the model has less than a certain threshold of BatchNorm layers.
Args:
thresh (int, optional): The threshold number of BatchNorm layers.
Returns:
(bool): True if the number of BatchNorm layers in the model is less than the threshold, False otherwise.
"""
bn = tuple(v for k, v in torch.nn.__dict__.items() if "Norm" in k) # normalization layers, i.e. BatchNorm2d()
return sum(isinstance(v, bn) for v in self.modules()) < thresh # True if < 'thresh' BatchNorm layers in model
def info(self, detailed=False, verbose=True, imgsz=640):
"""
Print model information.
Args:
detailed (bool): If True, prints out detailed information about the model.
verbose (bool): If True, prints out the model information.
imgsz (int): The size of the image that the model will be trained on.
"""
return model_info(self, detailed=detailed, verbose=verbose, imgsz=imgsz)
def _apply(self, fn):
"""
Apply a function to all tensors in the model that are not parameters or registered buffers.
Args:
fn (function): The function to apply to the model.
Returns:
(BaseModel): An updated BaseModel object.
"""
self = super()._apply(fn)
m = self.model[-1] # Detect()/Segment()/Pose()/OBB()
heads = DETECT_CLASS + SEGMENT_CLASS + POSE_CLASS + OBB_CLASS
if isinstance(m, heads):
m.stride = fn(m.stride)
m.anchors = fn(m.anchors)
m.strides = fn(m.strides)
return self
def load(self, weights, verbose=True):
"""
Load weights into the model.
Args:
weights (dict | torch.nn.Module): The pre-trained weights to be loaded.
verbose (bool, optional): Whether to log the transfer progress.
"""
model = weights["model"] if isinstance(weights, dict) else weights # torchvision models are not dicts
csd = model.float().state_dict() # checkpoint state_dict as FP32
updated_csd = intersect_dicts(csd, self.state_dict()) # intersect
self.load_state_dict(updated_csd, strict=False) # load
len_updated_csd = len(updated_csd)
first_conv = "model.0.conv.weight" # hard-coded to yolo models for now
# mostly used to boost multi-channel training
state_dict = self.state_dict()
if first_conv not in updated_csd and first_conv in state_dict:
c1, c2, h, w = state_dict[first_conv].shape
cc1, cc2, ch, cw = csd[first_conv].shape
if ch == h and cw == w:
c1, c2 = min(c1, cc1), min(c2, cc2)
state_dict[first_conv][:c1, :c2] = csd[first_conv][:c1, :c2]
len_updated_csd += 1
if verbose:
LOGGER.info(f"Transferred {len_updated_csd}/{len(self.model.state_dict())} items from pretrained weights")
def loss(self, batch, preds=None):
"""
Compute loss.
Args:
batch (dict): Batch to compute loss on.
preds (torch.Tensor | List[torch.Tensor], optional): Predictions.
"""
if getattr(self, "criterion", None) is None:
self.criterion = self.init_criterion()
preds = self.forward(batch["img"]) if preds is None else preds
det_loss_vec, det_items = self.criterion(preds, batch)
# Optional contrastive branch (enabled only if hooks exist and compute succeeds)
# Requirements: MULTIMODAL + CONTRAST + registered hooks via 6th field
mm_hm = getattr(self, 'mm_hook_manager', None)
has_hooks = False
if mm_hm is not None:
# Only enable contrast branch when YAML registered hooks (6th-field) exist
try:
has_hooks = bool(mm_hm.has_hooks())
except Exception:
has_hooks = False
use_contrast = (
self.training
and MULTIMODAL_AVAILABLE
and CONTRAST_AVAILABLE
and has_hooks
)
if not use_contrast:
return det_loss_vec, det_items
# Read configs from args with safe defaults
args = getattr(self, 'args', None)
cfg = ContrastConfig(
tau=getattr(args, 'contrast_tau', 0.07) if args is not None else 0.07,
proj_dim=getattr(args, 'contrast_dim', 128) if args is not None else 128,
lambda_weight=getattr(args, 'contrast_lambda', 0.1) if args is not None else 0.1,
max_rois_per_image=getattr(args, 'contrast_max_rois', 64) if args is not None else 64,
share_head=getattr(args, 'contrast_share_head', False) if args is not None else False,
preferred_stages=tuple(getattr(args, 'contrast_stages', ("P4", "P5", "P3"))) if args is not None else ("P4", "P5", "P3"),
)
# Lazy create controller (only when hooks exist)
if getattr(self, 'mm_contrast_controller', None) is None and use_contrast:
# Create on correct device to avoid CPU/CUDA mismatch when forward() is first called
try:
dev = next(self.parameters()).device
except StopIteration:
# Fallback to batch image device if model has no parameters (unlikely)
img = batch.get('img')
dev = img.device if isinstance(img, torch.Tensor) else torch.device('cpu')
self.mm_contrast_controller = ContrastController(cfg).to(dev)
# Collect hooked features (do not pop to allow external visualization; keep bounded by latest write)
hook_buffers = mm_hm.collect(pop=False)
loss_c, stats = self.mm_contrast_controller(hook_buffers, batch)
# Debug: detect non-finite contrastive loss (when enabled and computed)
try:
if loss_c is not None and not torch.isfinite(loss_c):
from ultralytics.utils import LOGGER as _LOGGER
_LOGGER.warning(f"[CL][loss] non-finite loss_c detected: {float(loss_c.detach().cpu())}")
except Exception:
pass
if loss_c is None:
return det_loss_vec, det_items # no valid pairs this step
# Compose outputs: scale contrast by lambda for backprop, but report raw value in items
lambda_w = cfg.lambda_weight
if det_loss_vec.dim() == 0:
# safety: ensure vector form
det_loss_vec = det_loss_vec.unsqueeze(0)
total_vec = torch.cat([det_loss_vec, (loss_c * lambda_w).unsqueeze(0)], dim=0)
total_items = torch.cat([det_items, loss_c.detach().unsqueeze(0)], dim=0)
# Optionally, expose stats via side-effect for loggers (trainer can read from model)
self._contrast_last_stats = stats
return total_vec, total_items
def init_criterion(self):
"""Initialize the loss criterion for the BaseModel."""
raise NotImplementedError("compute_loss() needs to be implemented by task heads")
class DetectionModel(BaseModel):
"""
YOLO detection model.
This class implements the YOLO detection architecture, handling model initialization, forward pass,
augmented inference, and loss computation for object detection tasks.
Attributes:
yaml (dict): Model configuration dictionary.
model (torch.nn.Sequential): The neural network model.
save (list): List of layer indices to save outputs from.
names (dict): Class names dictionary.
inplace (bool): Whether to use inplace operations.
end2end (bool): Whether the model uses end-to-end detection.
stride (torch.Tensor): Model stride values.
Methods:
__init__: Initialize the YOLO detection model.
_predict_augment: Perform augmented inference.
_descale_pred: De-scale predictions following augmented inference.
_clip_augmented: Clip YOLO augmented inference tails.
init_criterion: Initialize the loss criterion.
Examples:
Initialize a detection model
>>> model = DetectionModel("yolo11n.yaml", ch=3, nc=80)
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="yolo11n.yaml", ch=3, nc=None, verbose=True):
"""
Initialize the YOLO detection model with the given config and parameters.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
super().__init__()
self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg) # cfg dict
if self.yaml["backbone"][0][2] == "Silence":
LOGGER.warning(
"YOLOv9 `Silence` module is deprecated in favor of torch.nn.Identity. "
"Please delete local *.pt file and re-download the latest model checkpoint."
)
self.yaml["backbone"][0][2] = "nn.Identity"
# Define model
self.yaml["channels"] = ch # save channels
if nc and nc != self.yaml["nc"]:
LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
self.yaml["nc"] = nc # override YAML value
self.model, self.save = parse_model(deepcopy(self.yaml), ch=ch, verbose=verbose) # model, savelist
self.names = {i: f"{i}" for i in range(self.yaml["nc"])} # default names dict
self.inplace = self.yaml.get("inplace", True)
self.end2end = getattr(self.model[-1], "end2end", False)
# textmulti-modaltext(textexists)
if hasattr(self.model, 'multimodal_router'):
self.multimodal_router = self.model.multimodal_router
else:
self.multimodal_router = None
# Persist router for runtime ablation/filling so BaseModel forward can reuse it
self.mm_router = self.multimodal_router if self.multimodal_router is not None else None
# textcreate: textinfo()textLazytext
# actualcreatetexttrainingtextget_modeltext(textBaseModel.loss).
# textHookManager(textexists)
self.mm_hook_manager = getattr(self.model, 'mm_hook_manager', None)
# Build strides
m = self.model[-1] # Detect()/Segment()/Pose()/OBB()/...
heads = DETECT_CLASS + SEGMENT_CLASS + POSE_CLASS + OBB_CLASS
if isinstance(m, heads): # includes all Detect/Segment/Pose/OBB subclasses (e.g., LSCD variants)
s = 256 # 2x min stride
m.inplace = self.inplace
def _forward(x):
"""Perform a forward pass through the model, handling different Detect subclass types accordingly."""
if self.end2end:
return self.forward(x)["one2many"]
seg_pose_obb = SEGMENT_CLASS + POSE_CLASS + OBB_CLASS
return self.forward(x)[0] if isinstance(m, seg_pose_obb) else self.forward(x)
self.model.eval() # Avoid changing batch statistics until training begins
m.training = True # Setting it to True to properly return strides
m.stride = torch.tensor([s / x.shape[-2] for x in _forward(torch.zeros(1, ch, s, s))]) # forward
self.stride = m.stride
self.model.train() # Set model back to training(default) mode
# textdetectiontext(text)
if hasattr(m, "bias_init") and callable(getattr(m, "bias_init")):
m.bias_init() # only run once
else:
self.stride = torch.Tensor([32]) # default stride for i.e. RTDETR
# Init weights, biases
initialize_weights(self)
if verbose:
self.info()
LOGGER.info("")
def _predict_augment(self, x):
"""
Perform augmentations on input image x and return augmented inference and train outputs.
Args:
x (torch.Tensor): Input image tensor.
Returns:
(torch.Tensor): Augmented inference output.
"""
if getattr(self, "end2end", False) or self.__class__.__name__ != "DetectionModel":
LOGGER.warning("Model does not support 'augment=True', reverting to single-scale prediction.")
return self._predict_once(x)
img_size = x.shape[-2:] # height, width
s = [1, 0.83, 0.67] # scales
f = [None, 3, None] # flips (2-ud, 3-lr)
y = [] # outputs
for si, fi in zip(s, f):
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
yi = super().predict(xi)[0] # forward
yi = self._descale_pred(yi, fi, si, img_size)
y.append(yi)
y = self._clip_augmented(y) # clip augmented tails
return torch.cat(y, -1), None # augmented inference, train
@staticmethod
def _descale_pred(p, flips, scale, img_size, dim=1):
"""
De-scale predictions following augmented inference (inverse operation).
Args:
p (torch.Tensor): Predictions tensor.
flips (int): Flip type (0=none, 2=ud, 3=lr).
scale (float): Scale factor.
img_size (tuple): Original image size (height, width).
dim (int): Dimension to split at.
Returns:
(torch.Tensor): De-scaled predictions.
"""
p[:, :4] /= scale # de-scale
x, y, wh, cls = p.split((1, 1, 2, p.shape[dim] - 4), dim)
if flips == 2:
y = img_size[0] - y # de-flip ud
elif flips == 3:
x = img_size[1] - x # de-flip lr
return torch.cat((x, y, wh, cls), dim)
def _clip_augmented(self, y):
"""
Clip YOLO augmented inference tails.
Args:
y (List[torch.Tensor]): List of detection tensors.
Returns:
(List[torch.Tensor]): Clipped detection tensors.
"""
nl = self.model[-1].nl # number of detection layers (P3-P5)
g = sum(4**x for x in range(nl)) # grid points
e = 1 # exclude layer count
i = (y[0].shape[-1] // g) * sum(4**x for x in range(e)) # indices
y[0] = y[0][..., :-i] # large
i = (y[-1].shape[-1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # indices
y[-1] = y[-1][..., i:] # small
return y
def init_criterion(self):
"""Initialize the loss criterion for the DetectionModel."""
return E2EDetectLoss(self) if getattr(self, "end2end", False) else v8DetectionLoss(self)
class OBBModel(DetectionModel):
"""
YOLO Oriented Bounding Box (OBB) model.
This class extends DetectionModel to handle oriented bounding box detection tasks, providing specialized
loss computation for rotated object detection.
Methods:
__init__: Initialize YOLO OBB model.
init_criterion: Initialize the loss criterion for OBB detection.
Examples:
Initialize an OBB model
>>> model = OBBModel("yolo11n-obb.yaml", ch=3, nc=80)
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="yolo11n-obb.yaml", ch=3, nc=None, verbose=True):
"""
Initialize YOLO OBB model with given config and parameters.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
def init_criterion(self):
"""Initialize the loss criterion for the model."""
return v8OBBLoss(self)
class SegmentationModel(DetectionModel):
"""
YOLO segmentation model.
This class extends DetectionModel to handle instance segmentation tasks, providing specialized
loss computation for pixel-level object detection and segmentation.
Methods:
__init__: Initialize YOLO segmentation model.
init_criterion: Initialize the loss criterion for segmentation.
Examples:
Initialize a segmentation model
>>> model = SegmentationModel("yolo11n-seg.yaml", ch=3, nc=80)
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="yolo11n-seg.yaml", ch=3, nc=None, verbose=True):
"""
Initialize Ultralytics YOLO segmentation model with given config and parameters.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
def init_criterion(self):
"""Initialize the loss criterion for the SegmentationModel."""
return v8SegmentationLoss(self)
class PoseModel(DetectionModel):
"""
YOLO pose model.
This class extends DetectionModel to handle human pose estimation tasks, providing specialized
loss computation for keypoint detection and pose estimation.
Attributes:
kpt_shape (tuple): Shape of keypoints data (num_keypoints, num_dimensions).
Methods:
__init__: Initialize YOLO pose model.
init_criterion: Initialize the loss criterion for pose estimation.
Examples:
Initialize a pose model
>>> model = PoseModel("yolo11n-pose.yaml", ch=3, nc=1, data_kpt_shape=(17, 3))
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="yolo11n-pose.yaml", ch=3, nc=None, data_kpt_shape=(None, None), verbose=True):
"""
Initialize Ultralytics YOLO Pose model.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
data_kpt_shape (tuple): Shape of keypoints data.
verbose (bool): Whether to display model information.
"""
if not isinstance(cfg, dict):
cfg = yaml_model_load(cfg) # load model YAML
if any(data_kpt_shape) and list(data_kpt_shape) != list(cfg["kpt_shape"]):
LOGGER.info(f"Overriding model.yaml kpt_shape={cfg['kpt_shape']} with kpt_shape={data_kpt_shape}")
cfg["kpt_shape"] = data_kpt_shape
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
def init_criterion(self):
"""Initialize the loss criterion for the PoseModel."""
return v8PoseLoss(self)
class ClassificationModel(BaseModel):
"""
YOLO classification model.
This class implements the YOLO classification architecture for image classification tasks,
providing model initialization, configuration, and output reshaping capabilities.
Attributes:
yaml (dict): Model configuration dictionary.
model (torch.nn.Sequential): The neural network model.
stride (torch.Tensor): Model stride values.
names (dict): Class names dictionary.
Methods:
__init__: Initialize ClassificationModel.
_from_yaml: Set model configurations and define architecture.
reshape_outputs: Update model to specified class count.
init_criterion: Initialize the loss criterion.
Examples:
Initialize a classification model
>>> model = ClassificationModel("yolo11n-cls.yaml", ch=3, nc=1000)
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="yolo11n-cls.yaml", ch=3, nc=None, verbose=True):
"""
Initialize ClassificationModel with YAML, channels, number of classes, verbose flag.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
super().__init__()
self._from_yaml(cfg, ch, nc, verbose)
def _from_yaml(self, cfg, ch, nc, verbose):
"""
Set Ultralytics YOLO model configurations and define the model architecture.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg) # cfg dict
# Define model
ch = self.yaml["channels"] = self.yaml.get("channels", ch) # input channels
if nc and nc != self.yaml["nc"]:
LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
self.yaml["nc"] = nc # override YAML value
elif not nc and not self.yaml.get("nc", None):
raise ValueError("nc not specified. Must specify nc in model.yaml or function arguments.")
self.model, self.save = parse_model(deepcopy(self.yaml), ch=ch, verbose=verbose) # model, savelist
self.stride = torch.Tensor([1]) # no stride constraints
self.names = {i: f"{i}" for i in range(self.yaml["nc"])} # default names dict
self.info()
@staticmethod
def reshape_outputs(model, nc):
"""
Update a TorchVision classification model to class count 'n' if required.
Args:
model (torch.nn.Module): Model to update.
nc (int): New number of classes.
"""
name, m = list((model.model if hasattr(model, "model") else model).named_children())[-1] # last module
if isinstance(m, Classify): # YOLO Classify() head
if m.linear.out_features != nc:
m.linear = torch.nn.Linear(m.linear.in_features, nc)
elif isinstance(m, torch.nn.Linear): # ResNet, EfficientNet
if m.out_features != nc:
setattr(model, name, torch.nn.Linear(m.in_features, nc))
elif isinstance(m, torch.nn.Sequential):
types = [type(x) for x in m]
if torch.nn.Linear in types:
i = len(types) - 1 - types[::-1].index(torch.nn.Linear) # last torch.nn.Linear index
if m[i].out_features != nc:
m[i] = torch.nn.Linear(m[i].in_features, nc)
elif torch.nn.Conv2d in types:
i = len(types) - 1 - types[::-1].index(torch.nn.Conv2d) # last torch.nn.Conv2d index
if m[i].out_channels != nc:
m[i] = torch.nn.Conv2d(
m[i].in_channels, nc, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None
)
def init_criterion(self):
"""Initialize the loss criterion for the ClassificationModel."""
return v8ClassificationLoss()
class RTDETRDetectionModel(DetectionModel):
"""
RTDETR (Real-time DEtection and Tracking using Transformers) Detection Model class.
This class is responsible for constructing the RTDETR architecture, defining loss functions, and facilitating both
the training and inference processes. RTDETR is an object detection and tracking model that extends from the
DetectionModel base class.
Attributes:
nc (int): Number of classes for detection.
criterion (RTDETRDetectionLoss): Loss function for training.
Methods:
__init__: Initialize the RTDETRDetectionModel.
init_criterion: Initialize the loss criterion.
loss: Compute loss for training.
predict: Perform forward pass through the model.
Examples:
Initialize an RTDETR model
>>> model = RTDETRDetectionModel("rtdetr-l.yaml", ch=3, nc=80)
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="rtdetr-l.yaml", ch=3, nc=None, verbose=True):
"""
Initialize the RTDETRDetectionModel.
Args:
cfg (str | dict): Configuration file name or path.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Print additional information during initialization.
"""
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
# ===== MULTIMODAL EXTENSION START - textMultiModalRouter =====
self.mm_router = None
self.mm_routing_enabled = False
self.mm_input_sources = None
if MULTIMODAL_AVAILABLE:
try:
from ultralytics.nn.mm import MultiModalRouter
# Create persistent router with model configuration
config_dict = getattr(self, 'yaml', None)
self.mm_router = MultiModalRouter(config_dict, verbose=verbose)
if verbose:
LOGGER.info("RTDETRDetectionModel: textMultiModalRoutercreate")
except Exception as e:
if verbose:
LOGGER.warning(f"RTDETRDetectionModel: MultiModalRoutertextfailed: {e}")
# ===== MULTIMODAL EXTENSION END =====
def init_criterion(self):
"""Initialize the loss criterion for the RTDETRDetectionModel."""
from ultralytics.models.utils.loss import RTDETRDetectionLoss
return RTDETRDetectionLoss(nc=self.nc, use_vfl=True)
def loss(self, batch, preds=None):
"""
Compute the loss for the given batch of data.
Args:
batch (dict): Dictionary containing image and label data.
preds (torch.Tensor, optional): Precomputed model predictions.
Returns:
loss_sum (torch.Tensor): Total loss value.
loss_items (torch.Tensor): Main three losses in a tensor.
"""
if not hasattr(self, "criterion"):
self.criterion = self.init_criterion()
img = batch["img"]
# NOTE: preprocess gt_bbox and gt_labels to list.
bs = len(img)
batch_idx = batch["batch_idx"]
gt_groups = [(batch_idx == i).sum().item() for i in range(bs)]
targets = {
"cls": batch["cls"].to(img.device, dtype=torch.long).view(-1),
"bboxes": batch["bboxes"].to(device=img.device),
"batch_idx": batch_idx.to(img.device, dtype=torch.long).view(-1),
"gt_groups": gt_groups,
}
preds = self.predict(img, batch=targets) if preds is None else preds
dec_bboxes, dec_scores, enc_bboxes, enc_scores, dn_meta = preds if self.training else preds[1]
if dn_meta is None:
dn_bboxes, dn_scores = None, None
else:
dn_bboxes, dec_bboxes = torch.split(dec_bboxes, dn_meta["dn_num_split"], dim=2)
dn_scores, dec_scores = torch.split(dec_scores, dn_meta["dn_num_split"], dim=2)
dec_bboxes = torch.cat([enc_bboxes.unsqueeze(0), dec_bboxes]) # (7, bs, 300, 4)
dec_scores = torch.cat([enc_scores.unsqueeze(0), dec_scores])
loss = self.criterion(
(dec_bboxes, dec_scores), targets, dn_bboxes=dn_bboxes, dn_scores=dn_scores, dn_meta=dn_meta
)
# NOTE: There are like 12 losses in RTDETR, backward with all losses but only show the main three losses.
return sum(loss.values()), torch.as_tensor(
[loss[k].detach() for k in ["loss_giou", "loss_class", "loss_bbox"]], device=img.device
)
def predict(self, x, profile=False, visualize=False, batch=None, augment=False, embed=None):
"""
Perform a forward pass through the model.
Args:
x (torch.Tensor): The input tensor.
profile (bool): If True, profile the computation time for each layer.
visualize (bool): If True, save feature maps for visualization.
batch (dict, optional): Ground truth data for evaluation.
augment (bool): If True, perform data augmentation during inference.
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(torch.Tensor): Model's output tensor.
"""
# ===== MULTIMODAL EXTENSION START - multi-modaltext =====
mm_router = self.mm_router
mm_routing_enabled, mm_input_sources = (mm_router.setup_multimodal_routing(x, profile) if mm_router is not None else (False, None))
# ===== MULTIMODAL EXTENSION END =====
y, dt, embeddings = [], [], [] # outputs
embed = frozenset(embed) if embed is not None else {-1}
max_idx = max(embed)
for m in self.model[:-1]: # except the head part
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
# ===== MULTIMODAL EXTENSION START - multi-modaltext =====
# Apply multimodal routing if enabled and module has MM attributes
if mm_routing_enabled and mm_input_sources and mm_router:
routed_x = mm_router.route_layer_input(x, m, mm_input_sources, profile)
if routed_x is not None:
x = routed_x
# Check for spatial reset requirement
if mm_router and hasattr(m, '_mm_spatial_reset') and m._mm_spatial_reset:
x = mm_router.reset_spatial_input(x, m, mm_input_sources, profile)
# ===== MULTIMODAL EXTENSION END =====
if profile:
self._profile_one_layer(m, x, dt)
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
if m.i in embed:
embeddings.append(torch.nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1)) # flatten
if m.i == max_idx:
return torch.unbind(torch.cat(embeddings, 1), dim=0)
head = self.model[-1]
x = head([y[j] for j in head.f], batch) # head inference
return x
class WorldModel(DetectionModel):
"""
YOLOv8 World Model.
This class implements the YOLOv8 World model for open-vocabulary object detection, supporting text-based
class specification and CLIP model integration for zero-shot detection capabilities.
Attributes:
txt_feats (torch.Tensor): Text feature embeddings for classes.
clip_model (torch.nn.Module): CLIP model for text encoding.
Methods:
__init__: Initialize YOLOv8 world model.
set_classes: Set classes for offline inference.
get_text_pe: Get text positional embeddings.
predict: Perform forward pass with text features.
loss: Compute loss with text features.
Examples:
Initialize a world model
>>> model = WorldModel("yolov8s-world.yaml", ch=3, nc=80)
>>> model.set_classes(["person", "car", "bicycle"])
>>> results = model.predict(image_tensor)
"""
def __init__(self, cfg="yolov8s-world.yaml", ch=3, nc=None, verbose=True):
"""
Initialize YOLOv8 world model with given config and parameters.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
self.txt_feats = torch.randn(1, nc or 80, 512) # features placeholder
self.clip_model = None # CLIP model placeholder
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
def set_classes(self, text, batch=80, cache_clip_model=True):
"""
Set classes in advance so that model could do offline-inference without clip model.
Args:
text (List[str]): List of class names.
batch (int): Batch size for processing text tokens.
cache_clip_model (bool): Whether to cache the CLIP model.
"""
self.txt_feats = self.get_text_pe(text, batch=batch, cache_clip_model=cache_clip_model)
self.model[-1].nc = len(text)
def get_text_pe(self, text, batch=80, cache_clip_model=True):
"""
Set classes in advance so that model could do offline-inference without clip model.
Args:
text (List[str]): List of class names.
batch (int): Batch size for processing text tokens.
cache_clip_model (bool): Whether to cache the CLIP model.
Returns:
(torch.Tensor): Text positional embeddings.
"""
from ultralytics.nn.text_model import build_text_model
device = next(self.model.parameters()).device
if not getattr(self, "clip_model", None) and cache_clip_model:
# For backwards compatibility of models lacking clip_model attribute
self.clip_model = build_text_model("clip:ViT-B/32", device=device)
model = self.clip_model if cache_clip_model else build_text_model("clip:ViT-B/32", device=device)
text_token = model.tokenize(text)
txt_feats = [model.encode_text(token).detach() for token in text_token.split(batch)]
txt_feats = txt_feats[0] if len(txt_feats) == 1 else torch.cat(txt_feats, dim=0)
return txt_feats.reshape(-1, len(text), txt_feats.shape[-1])
def predict(self, x, profile=False, visualize=False, txt_feats=None, augment=False, embed=None):
"""
Perform a forward pass through the model.
Args:
x (torch.Tensor): The input tensor.
profile (bool): If True, profile the computation time for each layer.
visualize (bool): If True, save feature maps for visualization.
txt_feats (torch.Tensor, optional): The text features, use it if it's given.
augment (bool): If True, perform data augmentation during inference.
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(torch.Tensor): Model's output tensor.
"""
txt_feats = (self.txt_feats if txt_feats is None else txt_feats).to(device=x.device, dtype=x.dtype)
if len(txt_feats) != len(x) or self.model[-1].export:
txt_feats = txt_feats.expand(x.shape[0], -1, -1)
ori_txt_feats = txt_feats.clone()
y, dt, embeddings = [], [], [] # outputs
embed = frozenset(embed) if embed is not None else {-1}
max_idx = max(embed)
for m in self.model: # except the head part
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if profile:
self._profile_one_layer(m, x, dt)
if isinstance(m, C2fAttn):
x = m(x, txt_feats)
elif isinstance(m, WorldDetect):
x = m(x, ori_txt_feats)
elif isinstance(m, ImagePoolingAttn):
txt_feats = m(x, txt_feats)
else:
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
if m.i in embed:
embeddings.append(torch.nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1)) # flatten
if m.i == max_idx:
return torch.unbind(torch.cat(embeddings, 1), dim=0)
return x
def loss(self, batch, preds=None):
"""
Compute loss.
Args:
batch (dict): Batch to compute loss on.
preds (torch.Tensor | List[torch.Tensor], optional): Predictions.
"""
if not hasattr(self, "criterion"):
self.criterion = self.init_criterion()
if preds is None:
preds = self.forward(batch["img"], txt_feats=batch["txt_feats"])
return self.criterion(preds, batch)
class YOLOEModel(DetectionModel):
"""
YOLOE detection model.
This class implements the YOLOE architecture for efficient object detection with text and visual prompts,
supporting both prompt-based and prompt-free inference modes.
Attributes:
pe (torch.Tensor): Prompt embeddings for classes.
clip_model (torch.nn.Module): CLIP model for text encoding.
Methods:
__init__: Initialize YOLOE model.
get_text_pe: Get text positional embeddings.
get_visual_pe: Get visual embeddings.
set_vocab: Set vocabulary for prompt-free model.
get_vocab: Get fused vocabulary layer.
set_classes: Set classes for offline inference.
get_cls_pe: Get class positional embeddings.
predict: Perform forward pass with prompts.
loss: Compute loss with prompts.
Examples:
Initialize a YOLOE model
>>> model = YOLOEModel("yoloe-v8s.yaml", ch=3, nc=80)
>>> results = model.predict(image_tensor, tpe=text_embeddings)
"""
def __init__(self, cfg="yoloe-v8s.yaml", ch=3, nc=None, verbose=True):
"""
Initialize YOLOE model with given config and parameters.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
@smart_inference_mode()
def get_text_pe(self, text, batch=80, cache_clip_model=False, without_reprta=False):
"""
Set classes in advance so that model could do offline-inference without clip model.
Args:
text (List[str]): List of class names.
batch (int): Batch size for processing text tokens.
cache_clip_model (bool): Whether to cache the CLIP model.
without_reprta (bool): Whether to return text embeddings cooperated with reprta module.
Returns:
(torch.Tensor): Text positional embeddings.
"""
from ultralytics.nn.text_model import build_text_model
device = next(self.model.parameters()).device
if not getattr(self, "clip_model", None) and cache_clip_model:
# For backwards compatibility of models lacking clip_model attribute
self.clip_model = build_text_model("mobileclip:blt", device=device)
model = self.clip_model if cache_clip_model else build_text_model("mobileclip:blt", device=device)
text_token = model.tokenize(text)
txt_feats = [model.encode_text(token).detach() for token in text_token.split(batch)]
txt_feats = txt_feats[0] if len(txt_feats) == 1 else torch.cat(txt_feats, dim=0)
txt_feats = txt_feats.reshape(-1, len(text), txt_feats.shape[-1])
if without_reprta:
return txt_feats
assert not self.training
head = self.model[-1]
assert isinstance(head, YOLOEDetect)
return head.get_tpe(txt_feats) # run auxiliary text head
@smart_inference_mode()
def get_visual_pe(self, img, visual):
"""
Get visual embeddings.
Args:
img (torch.Tensor): Input image tensor.
visual (torch.Tensor): Visual features.
Returns:
(torch.Tensor): Visual positional embeddings.
"""
return self(img, vpe=visual, return_vpe=True)
def set_vocab(self, vocab, names):
"""
Set vocabulary for the prompt-free model.
Args:
vocab (nn.ModuleList): List of vocabulary items.
names (List[str]): List of class names.
"""
assert not self.training
head = self.model[-1]
assert isinstance(head, YOLOEDetect)
# Cache anchors for head
device = next(self.parameters()).device
self(torch.empty(1, 3, self.args["imgsz"], self.args["imgsz"]).to(device)) # warmup
# re-parameterization for prompt-free model
self.model[-1].lrpc = nn.ModuleList(
LRPCHead(cls, pf[-1], loc[-1], enabled=i != 2)
for i, (cls, pf, loc) in enumerate(zip(vocab, head.cv3, head.cv2))
)
for loc_head, cls_head in zip(head.cv2, head.cv3):
assert isinstance(loc_head, nn.Sequential)
assert isinstance(cls_head, nn.Sequential)
del loc_head[-1]
del cls_head[-1]
self.model[-1].nc = len(names)
self.names = check_class_names(names)
def get_vocab(self, names):
"""
Get fused vocabulary layer from the model.
Args:
names (list): List of class names.
Returns:
(nn.ModuleList): List of vocabulary modules.
"""
assert not self.training
head = self.model[-1]
assert isinstance(head, YOLOEDetect)
assert not head.is_fused
tpe = self.get_text_pe(names)
self.set_classes(names, tpe)
device = next(self.model.parameters()).device
head.fuse(self.pe.to(device)) # fuse prompt embeddings to classify head
vocab = nn.ModuleList()
for cls_head in head.cv3:
assert isinstance(cls_head, nn.Sequential)
vocab.append(cls_head[-1])
return vocab
def set_classes(self, names, embeddings):
"""
Set classes in advance so that model could do offline-inference without clip model.
Args:
names (List[str]): List of class names.
embeddings (torch.Tensor): Embeddings tensor.
"""
assert not hasattr(self.model[-1], "lrpc"), (
"Prompt-free model does not support setting classes. Please try with Text/Visual prompt models."
)
assert embeddings.ndim == 3
self.pe = embeddings
self.model[-1].nc = len(names)
self.names = check_class_names(names)
def get_cls_pe(self, tpe, vpe):
"""
Get class positional embeddings.
Args:
tpe (torch.Tensor, optional): Text positional embeddings.
vpe (torch.Tensor, optional): Visual positional embeddings.
Returns:
(torch.Tensor): Class positional embeddings.
"""
all_pe = []
if tpe is not None:
assert tpe.ndim == 3
all_pe.append(tpe)
if vpe is not None:
assert vpe.ndim == 3
all_pe.append(vpe)
if not all_pe:
all_pe.append(getattr(self, "pe", torch.zeros(1, 80, 512)))
return torch.cat(all_pe, dim=1)
def predict(
self, x, profile=False, visualize=False, tpe=None, augment=False, embed=None, vpe=None, return_vpe=False
):
"""
Perform a forward pass through the model.
Args:
x (torch.Tensor): The input tensor.
profile (bool): If True, profile the computation time for each layer.
visualize (bool): If True, save feature maps for visualization.
tpe (torch.Tensor, optional): Text positional embeddings.
augment (bool): If True, perform data augmentation during inference.
embed (list, optional): A list of feature vectors/embeddings to return.
vpe (torch.Tensor, optional): Visual positional embeddings.
return_vpe (bool): If True, return visual positional embeddings.
Returns:
(torch.Tensor): Model's output tensor.
"""
y, dt, embeddings = [], [], [] # outputs
b = x.shape[0]
embed = frozenset(embed) if embed is not None else {-1}
max_idx = max(embed)
for m in self.model: # except the head part
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if profile:
self._profile_one_layer(m, x, dt)
if isinstance(m, YOLOEDetect):
vpe = m.get_vpe(x, vpe) if vpe is not None else None
if return_vpe:
assert vpe is not None
assert not self.training
return vpe
cls_pe = self.get_cls_pe(m.get_tpe(tpe), vpe).to(device=x[0].device, dtype=x[0].dtype)
if cls_pe.shape[0] != b or m.export:
cls_pe = cls_pe.expand(b, -1, -1)
x = m(x, cls_pe)
else:
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
if m.i in embed:
embeddings.append(torch.nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1)) # flatten
if m.i == max_idx:
return torch.unbind(torch.cat(embeddings, 1), dim=0)
return x
def loss(self, batch, preds=None):
"""
Compute loss.
Args:
batch (dict): Batch to compute loss on.
preds (torch.Tensor | List[torch.Tensor], optional): Predictions.
"""
if not hasattr(self, "criterion"):
from ultralytics.utils.loss import TVPDetectLoss
visual_prompt = batch.get("visuals", None) is not None # TODO
self.criterion = TVPDetectLoss(self) if visual_prompt else self.init_criterion()
if preds is None:
preds = self.forward(batch["img"], tpe=batch.get("txt_feats", None), vpe=batch.get("visuals", None))
return self.criterion(preds, batch)
class YOLOESegModel(YOLOEModel, SegmentationModel):
"""
YOLOE segmentation model.
This class extends YOLOEModel to handle instance segmentation tasks with text and visual prompts,
providing specialized loss computation for pixel-level object detection and segmentation.
Methods:
__init__: Initialize YOLOE segmentation model.
loss: Compute loss with prompts for segmentation.
Examples:
Initialize a YOLOE segmentation model
>>> model = YOLOESegModel("yoloe-v8s-seg.yaml", ch=3, nc=80)
>>> results = model.predict(image_tensor, tpe=text_embeddings)
"""
def __init__(self, cfg="yoloe-v8s-seg.yaml", ch=3, nc=None, verbose=True):
"""
Initialize YOLOE segmentation model with given config and parameters.
Args:
cfg (str | dict): Model configuration file path or dictionary.
ch (int): Number of input channels.
nc (int, optional): Number of classes.
verbose (bool): Whether to display model information.
"""
super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
def loss(self, batch, preds=None):
"""
Compute loss.
Args:
batch (dict): Batch to compute loss on.
preds (torch.Tensor | List[torch.Tensor], optional): Predictions.
"""
if not hasattr(self, "criterion"):
from ultralytics.utils.loss import TVPSegmentLoss
visual_prompt = batch.get("visuals", None) is not None # TODO
self.criterion = TVPSegmentLoss(self) if visual_prompt else self.init_criterion()
if preds is None:
preds = self.forward(batch["img"], tpe=batch.get("txt_feats", None), vpe=batch.get("visuals", None))
return self.criterion(preds, batch)
class Ensemble(torch.nn.ModuleList):
"""
Ensemble of models.
This class allows combining multiple YOLO models into an ensemble for improved performance through
model averaging or other ensemble techniques.
Methods:
__init__: Initialize an ensemble of models.
forward: Generate predictions from all models in the ensemble.
Examples:
Create an ensemble of models
>>> ensemble = Ensemble()
>>> ensemble.append(model1)
>>> ensemble.append(model2)
>>> results = ensemble(image_tensor)
"""
def __init__(self):
"""Initialize an ensemble of models."""
super().__init__()
def forward(self, x, augment=False, profile=False, visualize=False):
"""
Generate the YOLO network's final layer.
Args:
x (torch.Tensor): Input tensor.
augment (bool): Whether to augment the input.
profile (bool): Whether to profile the model.
visualize (bool): Whether to visualize the features.
Returns:
y (torch.Tensor): Concatenated predictions from all models.
train_out (None): Always None for ensemble inference.
"""
y = [module(x, augment, profile, visualize)[0] for module in self]
# y = torch.stack(y).max(0)[0] # max ensemble
# y = torch.stack(y).mean(0) # mean ensemble
y = torch.cat(y, 2) # nms ensemble, y shape(B, HW, C)
return y, None # inference, train output
# Functions ------------------------------------------------------------------------------------------------------------
@contextlib.contextmanager
def temporary_modules(modules=None, attributes=None):
"""
Context manager for temporarily adding or modifying modules in Python's module cache (`sys.modules`).
This function can be used to change the module paths during runtime. It's useful when refactoring code,
where you've moved a module from one location to another, but you still want to support the old import
paths for backwards compatibility.
Args:
modules (dict, optional): A dictionary mapping old module paths to new module paths.
attributes (dict, optional): A dictionary mapping old module attributes to new module attributes.
Examples:
>>> with temporary_modules({"old.module": "new.module"}, {"old.module.attribute": "new.module.attribute"}):
>>> import old.module # this will now import new.module
>>> from old.module import attribute # this will now import new.module.attribute
Note:
The changes are only in effect inside the context manager and are undone once the context manager exits.
Be aware that directly manipulating `sys.modules` can lead to unpredictable results, especially in larger
applications or libraries. Use this function with caution.
"""
if modules is None:
modules = {}
if attributes is None:
attributes = {}
import sys
from importlib import import_module
try:
# Set attributes in sys.modules under their old name
for old, new in attributes.items():
old_module, old_attr = old.rsplit(".", 1)
new_module, new_attr = new.rsplit(".", 1)
setattr(import_module(old_module), old_attr, getattr(import_module(new_module), new_attr))
# Set modules in sys.modules under their old name
for old, new in modules.items():
sys.modules[old] = import_module(new)
yield
finally:
# Remove the temporary module paths
for old in modules:
if old in sys.modules:
del sys.modules[old]
class SafeClass:
"""A placeholder class to replace unknown classes during unpickling."""
def __init__(self, *args, **kwargs):
"""Initialize SafeClass instance, ignoring all arguments."""
pass
def __call__(self, *args, **kwargs):
"""Run SafeClass instance, ignoring all arguments."""
pass
class SafeUnpickler(pickle.Unpickler):
"""Custom Unpickler that replaces unknown classes with SafeClass."""
def find_class(self, module, name):
"""
Attempt to find a class, returning SafeClass if not among safe modules.
Args:
module (str): Module name.
name (str): Class name.
Returns:
(type): Found class or SafeClass.
"""
safe_modules = (
"torch",
"collections",
"collections.abc",
"builtins",
"math",
"numpy",
# Add other modules considered safe
)
if module in safe_modules:
return super().find_class(module, name)
else:
return SafeClass
def torch_safe_load(weight, safe_only=False):
"""
Attempt to load a PyTorch model with the torch.load() function. If a ModuleNotFoundError is raised, it catches the
error, logs a warning message, and attempts to install the missing module via the check_requirements() function.
After installation, the function again attempts to load the model using torch.load().
Args:
weight (str): The file path of the PyTorch model.
safe_only (bool): If True, replace unknown classes with SafeClass during loading.
Returns:
ckpt (dict): The loaded model checkpoint.
file (str): The loaded filename.
Examples:
>>> from ultralytics.nn.tasks import torch_safe_load
>>> ckpt, file = torch_safe_load("path/to/best.pt", safe_only=True)
"""
from ultralytics.utils.downloads import attempt_download_asset
check_suffix(file=weight, suffix=".pt")
file = attempt_download_asset(weight) # search online if missing locally
try:
with temporary_modules(
modules={
"ultralytics.yolo.utils": "ultralytics.utils",
"ultralytics.yolo.v8": "ultralytics.models.yolo",
"ultralytics.yolo.data": "ultralytics.data",
},
attributes={
"ultralytics.nn.modules.block.Silence": "torch.nn.Identity", # YOLOv9e
"ultralytics.nn.tasks.YOLOv10DetectionModel": "ultralytics.nn.tasks.DetectionModel", # YOLOv10
"ultralytics.utils.loss.v10DetectLoss": "ultralytics.utils.loss.E2EDetectLoss", # YOLOv10
},
):
if safe_only:
# Load via custom pickle module
safe_pickle = types.ModuleType("safe_pickle")
safe_pickle.Unpickler = SafeUnpickler
safe_pickle.load = lambda file_obj: SafeUnpickler(file_obj).load()
with open(file, "rb") as f:
ckpt = torch_load(f, pickle_module=safe_pickle)
else:
ckpt = torch_load(file, map_location="cpu")
except ModuleNotFoundError as e: # e.name is missing module name
if e.name == "models":
raise TypeError(
emojis(
f"ERROR ❌️ {weight} appears to be an Ultralytics YOLOv5 model originally trained "
f"with https://github.com/ultralytics/yolov5.\nThis model is NOT forwards compatible with "
f"YOLOv8 at https://github.com/ultralytics/ultralytics."
f"\nRecommend fixes are to train a new model using the latest 'ultralytics' package or to "
f"run a command with an official Ultralytics model, i.e. 'yolo predict model=yolo11n.pt'"
)
) from e
elif e.name == "numpy._core":
raise ModuleNotFoundError(
emojis(
f"ERROR ❌️ {weight} requires numpy>=1.26.1, however numpy=={__import__('numpy').__version__} is installed."
)
) from e
LOGGER.warning(
f"{weight} appears to require '{e.name}', which is not in Ultralytics requirements."
f"\nAutoInstall will run now for '{e.name}' but this feature will be removed in the future."
f"\nRecommend fixes are to train a new model using the latest 'ultralytics' package or to "
f"run a command with an official Ultralytics model, i.e. 'yolo predict model=yolo11n.pt'"
)
check_requirements(e.name) # install missing module
ckpt = torch_load(file, map_location="cpu")
if not isinstance(ckpt, dict):
# File is likely a YOLO instance saved with i.e. torch.save(model, "saved_model.pt")
LOGGER.warning(
f"The file '{weight}' appears to be improperly saved or formatted. "
f"For optimal results, use model.save('filename.pt') to correctly save YOLO models."
)
ckpt = {"model": ckpt.model}
return ckpt, file
def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
"""
Load an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a.
Args:
weights (str | List[str]): Model weights path(s).
device (torch.device, optional): Device to load model to.
inplace (bool): Whether to do inplace operations.
fuse (bool): Whether to fuse model.
Returns:
(torch.nn.Module): Loaded model.
"""
ensemble = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
ckpt, w = torch_safe_load(w) # load ckpt
args = {**DEFAULT_CFG_DICT, **ckpt["train_args"]} if "train_args" in ckpt else None # combined args
model = (ckpt.get("ema") or ckpt["model"]).to(device).float() # FP32 model
# Model compatibility updates
model.args = args # attach args to model
model.pt_path = w # attach *.pt file path to model
model.task = getattr(model, "task", guess_model_task(model))
if not hasattr(model, "stride"):
model.stride = torch.tensor([32.0])
# Append
ensemble.append(model.fuse().eval() if fuse and hasattr(model, "fuse") else model.eval()) # model in eval mode
# Module updates
for m in ensemble.modules():
if hasattr(m, "inplace"):
m.inplace = inplace
elif isinstance(m, torch.nn.Upsample) and not hasattr(m, "recompute_scale_factor"):
m.recompute_scale_factor = None # torch 1.11.0 compatibility
# Return model
if len(ensemble) == 1:
return ensemble[-1]
# Return ensemble
LOGGER.info(f"Ensemble created with {weights}\n")
for k in "names", "nc", "yaml":
setattr(ensemble, k, getattr(ensemble[0], k))
ensemble.stride = ensemble[int(torch.argmax(torch.tensor([m.stride.max() for m in ensemble])))].stride
assert all(ensemble[0].nc == m.nc for m in ensemble), f"Models differ in class counts {[m.nc for m in ensemble]}"
return ensemble
def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
"""
Load a single model weights.
Args:
weight (str): Model weight path.
device (torch.device, optional): Device to load model to.
inplace (bool): Whether to do inplace operations.
fuse (bool): Whether to fuse model.
Returns:
model (torch.nn.Module): Loaded model.
ckpt (dict): Model checkpoint dictionary.
"""
ckpt, weight = torch_safe_load(weight) # load ckpt
args = {**DEFAULT_CFG_DICT, **(ckpt.get("train_args", {}))} # combine model and default args, preferring model args
model = (ckpt.get("ema") or ckpt["model"]).to(device).float() # FP32 model
# Model compatibility updates
model.args = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS} # attach args to model
model.pt_path = weight # attach *.pt file path to model
model.task = getattr(model, "task", guess_model_task(model))
if not hasattr(model, "stride"):
model.stride = torch.tensor([32.0])
model = model.fuse().eval() if fuse and hasattr(model, "fuse") else model.eval() # model in eval mode
# Module updates
for m in model.modules():
if hasattr(m, "inplace"):
m.inplace = inplace
elif isinstance(m, torch.nn.Upsample) and not hasattr(m, "recompute_scale_factor"):
m.recompute_scale_factor = None # torch 1.11.0 compatibility
# Return model and ckpt
return model, ckpt
def _validate_and_fill_dea_args(args, c_left):
"""Validate DEA args strictly and auto-fill only the channel whentextNone.
text: DEA(channel, kernel_size, p_kernel=None, m_kernel=None, reduction=16)
text: DEA([None, kernel_size]) text DEA([channel, kernel_size]).
textautotext, text.
"""
a = list(args) if isinstance(args, (list, tuple)) else [args]
# text [channel/None, kernel_size]
if len(a) == 2:
ch, ks = a
ch = c_left if ch is None else ch
if not isinstance(ch, int) or not isinstance(ks, int) or ks <= 0:
raise ValueError(
f"DEA expects [channel(or None), kernel_size] as minimal form, got {args}"
)
return [ch, ks, None, None, 16]
# text: text5
while len(a) < 5:
a.append(None)
ch, ks, pk, mk, rd = a[:5]
ch = c_left if ch is None else ch
# text
if not isinstance(ch, int) or ch <= 0:
raise ValueError(f"DEA arg[0]=channel must be positive int, got {ch}")
if not isinstance(ks, int) or ks <= 0:
raise ValueError(f"DEA arg[1]=kernel_size must be positive int, got {ks}")
if pk is not None and (not isinstance(pk, (list, tuple)) or len(pk) != 2):
raise ValueError(f"DEA arg[2]=p_kernel must be 2-list/tuple or None, got {pk}")
if mk is not None and (not isinstance(mk, (list, tuple)) or len(mk) != 2):
raise ValueError(f"DEA arg[3]=m_kernel must be 2-list/tuple or None, got {mk}")
if rd is None:
rd = 16
if not isinstance(rd, int) or rd <= 0:
raise ValueError(f"DEA arg[4]=reduction must be positive int, got {rd}")
return [ch, ks, pk, mk, rd]
def parse_model(d, ch, verbose=True, dataset_config=None):
"""Parse the NeuroSeg-MF RT-DETR YAML into a PyTorch model.
Supported modules are intentionally limited to the current NeuroSeg-MF architecture:
ConvNormLayer, BasicBlock, Blocks, GateFusion, A2C2f, Conv, AIFI, Concat,
RepC3, RTDETRDecoder, nn.MaxPool2d and nn.Upsample.
"""
import ast
max_channels = float("inf")
nc, act, scales = (d.get(x) for x in ("nc", "activation", "scales"))
depth, width = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple"))
if scales:
scale = d.get("scale") or tuple(scales.keys())[0]
if not d.get("scale"):
LOGGER.warning(f"no model scale passed. Assuming scale='{scale}'.")
depth, width, max_channels = scales[scale]
if act:
allowed_acts = {
"torch.nn.SiLU()": torch.nn.SiLU(),
"nn.SiLU()": torch.nn.SiLU(),
"torch.nn.ReLU()": torch.nn.ReLU(),
"nn.ReLU()": torch.nn.ReLU(),
}
if act not in allowed_acts:
raise ValueError(f"Unsupported activation expression in minimal parser: {act}")
Conv.default_act = allowed_acts[act]
if verbose:
LOGGER.info(f"{colorstr('activation:')} {act}")
if verbose:
LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10} {'module':<45}{'arguments':<30}")
mm_router = None
if MULTIMODAL_AVAILABLE:
try:
config_dict = d.copy()
if dataset_config:
config_dict['dataset_config'] = dataset_config
mm_router = MultiModalRouter(config_dict, verbose=verbose)
except Exception as e:
if verbose:
LOGGER.warning(f"MultiModal router initialization failed: {e}")
ch = [ch]
layers, save = [], []
base_modules = {Conv, ConvNormLayer, A2C2f}
repeat_modules = {A2C2f, RepC3}
for i, layer_config in enumerate(d["backbone"] + d["head"]):
if mm_router:
c1, mm_input_source, mm_attributes = mm_router.parse_layer_config(layer_config, i, ch, verbose)
f, n, m, args = layer_config[:4]
else:
if len(layer_config) < 4:
raise ValueError(f"Invalid layer definition at index {i}: {layer_config}")
f, n, m, args = layer_config[:4]
c1, mm_input_source, mm_attributes = None, None, {}
m = getattr(torch.nn, m[3:]) if isinstance(m, str) and m.startswith("nn.") else globals().get(m)
if m is None:
raise ImportError(f"Module '{layer_config[2]}' used at layer {i} is not available in the NeuroSeg-MF minimal runtime.")
args = list(args)
for j, a in enumerate(args):
if isinstance(a, str):
if a == "nc":
args[j] = nc
elif a in globals():
args[j] = globals()[a]
else:
with contextlib.suppress(Exception):
args[j] = ast.literal_eval(a)
n_ = max(round(n * depth), 1) if n > 1 else n
n = n_
if m in base_modules:
if mm_input_source and c1 is not None:
c2 = args[0]
else:
c1, c2 = ch[f], args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels) * width, 8)
args = [c1, c2, *args[1:]]
if m in repeat_modules:
args.insert(2, n)
n = 1
if m is A2C2f:
args.extend((True, 1.2))
elif m is Blocks:
block_type = globals()[args[1]] if isinstance(args[1], str) else args[1]
c1, c2 = ch[f], args[0] * block_type.expansion
args = [c1, args[0], block_type, *args[2:]]
elif m is BasicBlock:
c1, c2 = ch[f], args[0] * BasicBlock.expansion
args = [c1, *args]
elif m is RepC3:
c1, c2 = ch[f], args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels) * width, 8)
args = [c1, c2, n, *args[1:]]
n = 1
elif m is AIFI:
c2 = ch[f]
args = [ch[f], *args]
elif m is Concat:
c2 = sum(ch[x] for x in f)
elif m is GateFusion:
c2 = ch[f[0]]
elif m is RTDETRDecoder:
args.insert(1, [ch[x] for x in f])
c2 = nc
elif m is torch.nn.BatchNorm2d:
args = [ch[f]]
c2 = ch[f]
elif m in {torch.nn.MaxPool2d, torch.nn.Upsample, torch.nn.Identity}:
c2 = ch[f]
elif m is Index:
c2 = ch[f]
else:
raise ImportError(f"Module '{layer_config[2]}' at layer {i} is not supported by the NeuroSeg-MF minimal parser.")
m_ = torch.nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)
t = str(m)[8:-2].replace("__main__.", "")
m_.np = sum(x.numel() for x in m_.parameters())
m_.i, m_.f, m_.type = i, f, t
if mm_attributes and mm_router:
mm_router.set_module_attributes(m_, mm_attributes)
if verbose:
display_args = [a.__name__ if isinstance(a, type) else a for a in args]
LOGGER.info(f"{i:>3}{str(f):>20}{n_:>3}{m_.np:10.0f} {t:<45}{str(display_args):<30}")
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)
layers.append(m_)
if i == 0:
ch = []
ch.append(c2)
model = torch.nn.Sequential(*layers)
if mm_router:
model.multimodal_router = mm_router
return model, sorted(save)
def yaml_model_load(path):
"""
Load a YOLOv8 model from a YAML file.
Args:
path (str | Path): Path to the YAML file.
Returns:
(dict): Model dictionary.
"""
path = Path(path)
if path.stem in (f"yolov{d}{x}6" for x in "nsmlx" for d in (5, 8)):
new_stem = re.sub(r"(\d+)([nslmx])6(.+)?$", r"\1\2-p6\3", path.stem)
LOGGER.warning(f"Ultralytics YOLO P6 models now use -p6 suffix. Renaming {path.stem} to {new_stem}.")
path = path.with_name(new_stem + path.suffix)
unified_path = re.sub(r"(\d+)([nslmx])(.+)?$", r"\1\3", str(path)) # i.e. yolov8x.yaml -> yolov8.yaml
yaml_file = check_yaml(unified_path, hard=False) or check_yaml(path)
d = YAML.load(yaml_file) # model dict
d["scale"] = guess_model_scale(path)
d["yaml_file"] = str(path)
return d
def guess_model_scale(model_path):
"""
Extract the size character n, s, m, l, or x of the model's scale from the model path.
Args:
model_path (str | Path): The path to the YOLO model's YAML file.
Returns:
(str): The size character of the model's scale (n, s, m, l, or x).
"""
try:
return re.search(r"yolo(e-)?[v]?\d+([nslmx])", Path(model_path).stem).group(2) # noqa
except AttributeError:
return ""
def guess_model_task(model):
"""
Guess the task of a PyTorch model from its architecture or configuration.
Args:
model (torch.nn.Module | dict): PyTorch model or model configuration in YAML format.
Returns:
(str): Task of the model ('detect', 'segment', 'classify', 'pose', 'obb').
"""
def cfg2task(cfg):
"""Guess from YAML dictionary."""
m = cfg["head"][-1][-2].lower() # output module name
if m in {"classify", "classifier", "cls", "fc"}:
return "classify"
if "detect" in m:
return "detect"
if "segment" in m:
return "segment"
if m == "pose":
return "pose"
if m == "obb":
return "obb"
# Guess from model cfg
if isinstance(model, dict):
with contextlib.suppress(Exception):
return cfg2task(model)
# Guess from PyTorch model
if isinstance(model, torch.nn.Module): # PyTorch model
def _resolve_attr(obj, dotted):
for name in dotted.split("."):
obj = getattr(obj, name)
return obj
for x in "model.args", "model.model.args", "model.model.model.args":
with contextlib.suppress(Exception):
return _resolve_attr(model, x)["task"]
for x in "model.yaml", "model.model.yaml", "model.model.model.yaml":
with contextlib.suppress(Exception):
return cfg2task(_resolve_attr(model, x))
for m in model.modules():
if isinstance(m, (Segment, YOLOESegment)):
return "segment"
elif isinstance(m, Classify):
return "classify"
elif isinstance(m, Pose):
return "pose"
elif isinstance(m, OBB):
return "obb"
elif isinstance(m, (Detect, WorldDetect, YOLOEDetect, v10Detect)):
return "detect"
# Guess from model filename
if isinstance(model, (str, Path)):
model = Path(model)
if "-seg" in model.stem or "segment" in model.parts:
return "segment"
elif "-cls" in model.stem or "classify" in model.parts:
return "classify"
elif "-pose" in model.stem or "pose" in model.parts:
return "pose"
elif "-obb" in model.stem or "obb" in model.parts:
return "obb"
elif "detect" in model.parts:
return "detect"
# Unable to determine task from model
LOGGER.warning(
"Unable to automatically guess model task, assuming 'task=detect'. "
"Explicitly define task for your model, i.e. 'task=detect', 'segment', 'classify','pose' or 'obb'."
)
return "detect" # assume detect

tasks.py at commit 4bfb1ef, no license · at the source

Overview

Authors: Zhehao Xu¹, Weiyi Liu², Shanshan Liang², Hongbo Jia^3,4, Xiaowei Chen^2,5, Han Qin⁵, Xiang Liao¹

ORCID iDs: Xiang Liao

Center for Neurointelligence, School of Medicine, Chongqing University, Chongqing 400030, China
Brain Research Center, State Key Laboratory of Trauma and Chemical Poisoning, Third Military Medical University, Chongqing 400038, China
Jiangsu Key Laboratory for Advanced Theranostics and Medical Instrumentation, Suzhou Institute of Biomedical Engineering and Technology, Chinese Academy of Sciences, Suzhou 215163, Jiangsu, China
Leibniz Institute for Neurobiology, Magdeburg 39118, Germany
LFC Laboratory (Chongqing Key Laboratory of Brain and Aerospace Intelligence) and Chongqing Institute for Brain and Intelligence, Guangyang Bay Laboratory, Chongqing 400064, China

Institutions: Chongqing University (China); Army Medical University (China); Chinese Academy of Sciences (China); Suzhou Institute of Biomedical Engineering and Technology (China); Leibniz Institute for Neurobiology (Germany)

Journal: Biomedical optics express, volume 17, issue 7, pages 3727-3746

Dates: received 31 March 2026; accepted 9 June 2026; published online 17 June 2026

Type: Research article · Language: English

License: none stated

Identifiers: DOI 10.1364/boe.600665 · PMID 42460356 · PMCID PMC13372336 · OpenAlex W7164500274

Open access: gold, a free copy (OpenAlex)

Status: code verified

Categories: optical imaging (calcium, voltage, 2-photon) (modality), methods / tools (subfield)

Methods: Statistics, Machine learning, Connectivity, fMRI & imaging

Topic: Advanced Fluorescence Microscopy Techniques (Biophysics, Biochemistry, Genetics and Molecular Biology), according to OpenAlex

Funding: National Natural Science Foundation of China (32171096, 32127801)

Citations: cited by 1 paper (Europe PMC); 57 references in the paper

Abstract

The abstract is not reproduced here: the paper's license (none stated) does not allow it. Read it in the paper, at the publisher or on Europe PMC.

Repository

Its files are read in the Code ↔ Paper reader above, with 11 matches between paragraphs and lines of code.

XZH-James/NeuroSeg-MF

License: none: the authors keep all their rights

State: the link answers, verified on 27 September 2026

Evidence: files inventoried

Commit: 4bfb1ef734acd2fa90b44960980d6233d71ed7cb, 16 June 2026

Languages: Python (100)

Size: 108 files, 100 scripts

Software Heritage: not archived

Found in: the references

Holds: README, environment (requirements.txt)

Not found: license file, CITATION.cff, tests, continuous integration, documentation

Tools: PyTorch (56 files), NumPy (39 files), OpenCV (20 files), Pillow (10 files), Matplotlib (5 files), pandas (4 files), SciPy (4 files), tifffile (2 files), seaborn (1 file), TensorFlow (1 file)

Availability: 1 check, the latest on 27 September 2026: the link answers

27 September 2026: the link answers

101 files

depth_anything_v2/dinov2.py, Python, 415 lines
depth_anything_v2/dinov2_layers/__init__.py, Python, 11 lines
depth_anything_v2/dinov2_layers/attention.py, Python, 83 lines
depth_anything_v2/dinov2_layers/block.py, Python, 252 lines
depth_anything_v2/dinov2_layers/drop_path.py, Python, 35 lines
depth_anything_v2/dinov2_layers/layer_scale.py, Python, 28 lines
depth_anything_v2/dinov2_layers/mlp.py, Python, 41 lines
depth_anything_v2/dinov2_layers/patch_embed.py, Python, 89 lines
depth_anything_v2/dinov2_layers/swiglu_ffn.py, Python, 63 lines
depth_anything_v2/dpt.py, Python, 221 lines
depth_anything_v2/util/blocks.py, Python, 148 lines
depth_anything_v2/util/transform.py, Python, 158 lines
predict.py, Python, 31 lines
tools/SAM_box_to_mask.py, Python, 675 lines, 2 matches
tools/generate_corr.py, Python, 591 lines, 1 match
tools/generate_depth_twophoton.py, Python, 755 lines, 2 matches
train.py, Python, 11 lines
ultralytics/__init__.py, Python, 14 lines
ultralytics/cfg/__init__.py, Python, 933 lines
ultralytics/data/__init__.py, Python, 61 lines
ultralytics/data/annotator.py, Python, 67 lines
ultralytics/data/augment.py, Python, 2,991 lines, 1 match
ultralytics/data/base.py, Python, 448 lines
ultralytics/data/build.py, Python, 426 lines
ultralytics/data/converter.py, Python, 756 lines
ultralytics/data/dataset.py, Python, 1,490 lines
ultralytics/data/loaders.py, Python, 708 lines
ultralytics/data/multimodal_augment.py, Python, 1,357 lines
ultralytics/data/split.py, Python, 138 lines
ultralytics/data/utils.py, Python, 796 lines
ultralytics/engine/__init__.py, Python, 1 line
ultralytics/engine/exporter.py, Python, 40 lines
ultralytics/engine/model.py, Python, 1,319 lines
ultralytics/engine/predictor.py, Python, 520 lines
ultralytics/engine/results.py, Python, 1,663 lines
ultralytics/engine/trainer.py, Python, 874 lines
ultralytics/engine/tuner.py, Python, 246 lines
ultralytics/engine/validator.py, Python, 370 lines
ultralytics/hub/__init__.py, Python, 7 lines
ultralytics/models/__init__.py, Python, 3 lines
ultralytics/models/rtdetr/__init__.py, Python, 13 lines
ultralytics/models/rtdetr/model.py, Python, 352 lines
ultralytics/models/rtdetr/multimodal/__init__.py, Python, 13 lines
ultralytics/models/rtdetr/multimodal/predict.py, Python, 1,204 lines
ultralytics/models/rtdetr/multimodal/train.py, Python, 568 lines
ultralytics/models/rtdetr/multimodal/val.py, Python, 845 lines
ultralytics/models/rtdetr/predict.py, Python, 91 lines
ultralytics/models/rtdetr/train.py, Python, 91 lines
ultralytics/models/rtdetr/val.py, Python, 231 lines
ultralytics/models/utils/__init__.py, Python, 1 line
ultralytics/models/utils/loss.py, Python, 476 lines
ultralytics/models/utils/multimodal/__init__.py, Python, 1 line
ultralytics/models/utils/multimodal/vis.py, Python, 35 lines
ultralytics/models/utils/ops.py, Python, 317 lines
ultralytics/models/yolo/__init__.py, Python, 2 lines
ultralytics/models/yolo/detect/__init__.py, Python, 3 lines
ultralytics/models/yolo/detect/train.py, Python, 58 lines
ultralytics/models/yolo/detect/val.py, Python, 75 lines
ultralytics/nn/__init__.py, Python, 29 lines
ultralytics/nn/autobackend.py, Python, 896 lines
ultralytics/nn/mm/__init__.py, Python, 68 lines
ultralytics/nn/mm/filling.py, Python, 174 lines
ultralytics/nn/mm/parser.py, Python, 195 lines
ultralytics/nn/mm/router.py, Python, 505 lines
ultralytics/nn/mm/utils.py, Python, 92 lines
ultralytics/nn/modules/__init__.py, Python, 71 lines, 2 matches
ultralytics/nn/modules/block.py, Python, 663 lines
ultralytics/nn/modules/conv.py, Python, 416 lines
ultralytics/nn/modules/head.py, Python, 341 lines
ultralytics/nn/modules/transformer.py, Python, 639 lines
ultralytics/nn/modules/utils.py, Python, 164 lines
ultralytics/nn/tasks.py, Python, 2,035 lines, 2 matches
ultralytics/nn/text_model.py, Python, 4 lines
ultralytics/solutions/__init__.py, Python, 1 line
ultralytics/solutions/config.py, Python, 7 lines
ultralytics/utils/__init__.py, Python, 1,599 lines
ultralytics/utils/autobatch.py, Python, 119 lines
ultralytics/utils/autodevice.py, Python, 206 lines
ultralytics/utils/benchmarks.py, Python, 720 lines
ultralytics/utils/callbacks/__init__.py, Python, 1 line
ultralytics/utils/callbacks/base.py, Python, 234 lines
ultralytics/utils/checks.py, Python, 936 lines
ultralytics/utils/coco_eval_bbox_mm.py, Python, 440 lines
ultralytics/utils/coco_metrics.py, Python, 1,733 lines, 1 match
ultralytics/utils/dist.py, Python, 119 lines
ultralytics/utils/downloads.py, Python, 529 lines
ultralytics/utils/errors.py, Python, 43 lines
ultralytics/utils/export.py, Python, 236 lines
ultralytics/utils/files.py, Python, 222 lines
ultralytics/utils/instance.py, Python, 504 lines
ultralytics/utils/loss.py, Python, 863 lines
ultralytics/utils/metrics.py, Python, 1,466 lines
ultralytics/utils/ops.py, Python, 888 lines
ultralytics/utils/patches.py, Python, 187 lines
ultralytics/utils/plotting.py, Python, 1,126 lines
ultralytics/utils/tal.py, Python, 419 lines
ultralytics/utils/torch_utils.py, Python, 1,064 lines
ultralytics/utils/triton.py, Python, 118 lines
ultralytics/utils/tuner.py, Python, 159 lines
val.py, Python, 8 lines
README.md, Text, 482 lines

At the source: github.com/XZH-James/NeuroSeg-MF

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
100 scripts, each with its path and the digest of its content;
11 matches between paragraphs of the paper and lines of the code (method lexical-v1);
neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Code and data availability statement

The paper has a code and data availability statement. Its license (none stated) does not allow reproducing it here; in short, from what the harvester recognized in it:

it says that the data are available on request

Read it in the paper: doi.org/10.1364/boe.600665.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 27 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 7 authors, 1 funder, 45 references.

Cite

This paper

Xu, Z., Liu, W., Liang, S., Jia, H., Chen, X., Qin, H., & Liao, X. (2026). NeuroSeg-MF: robust neuron segmentation in two-photon Ca2+ imaging using multi-feature fusion and detection-guided SAM. Biomedical optics express, 17(7), 3727-3746. https://doi.org/10.1364/boe.600665

BibTeX

@article{xu2026neuroseg, author = {Xu, Zhehao and Liu, Weiyi and Liang, Shanshan and Jia, Hongbo and Chen, Xiaowei and Qin, Han and Liao, Xiang}, title = {{NeuroSeg-MF: robust neuron segmentation in two-photon Ca\<sup\>2+\</sup\> imaging using multi-feature fusion and detection-guided SAM}}, journal = {Biomedical optics express}, year = {2026}, month = jun, volume = {17}, number = {7}, pages = {3727--3746}, publisher = {Optica Publishing Group}, issn = {2156-7085}, doi = {10.1364/boe.600665}, url = {https://doi.org/10.1364/boe.600665}, pmid = {42460356}, pmcid = {PMC13372336} }

RIS

TY - JOUR AU - Xu, Zhehao AU - Liu, Weiyi AU - Liang, Shanshan AU - Jia, Hongbo AU - Chen, Xiaowei AU - Qin, Han AU - Liao, Xiang TI - NeuroSeg-MF: robust neuron segmentation in two-photon Ca2+ imaging using multi-feature fusion and detection-guided SAM T2 - Biomedical optics express J2 - Biomed Opt Express PY - 2026 DA - 2026/06/17 VL - 17 IS - 7 SP - 3727 EP - 3746 SN - 2156-7085 PB - Optica Publishing Group DO - 10.1364/boe.600665 UR - https://doi.org/10.1364/boe.600665 LA - en ER -

CSL-JSON

{ "id": "10.1364/boe.600665", "type": "article-journal", "title": "NeuroSeg-MF: robust neuron segmentation in two-photon Ca2+ imaging using multi-feature fusion and detection-guided SAM", "container-title": "Biomedical optics express", "author": [ { "family": "Xu", "given": "Zhehao" }, { "family": "Liu", "given": "Weiyi" }, { "family": "Liang", "given": "Shanshan" }, { "family": "Jia", "given": "Hongbo" }, { "family": "Chen", "given": "Xiaowei" }, { "family": "Qin", "given": "Han" }, { "family": "Liao", "given": "Xiang" } ], "container-title-short": "Biomed Opt Express", "volume": "17", "issue": "7", "page": "3727-3746", "DOI": "10.1364/boe.600665", "PMID": "42460356", "PMCID": "PMC13372336", "ISSN": "2156-7085", "publisher": "Optica Publishing Group", "URL": "https://doi.org/10.1364/boe.600665", "language": "en", "issued": { "date-parts": [ [ 2026, 6, 17 ] ] } }

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.21037/qims-2026-0792 [code]: An nnU-Net-based framework with adaptive feature representation for 3D brain tumor segmentation.
Journal: Quantitative imaging in medicine and surgery
In common: tifffile, OpenCV, PyTorch, 5 other tools, methods / tools, 5 references
[2] doi:10.1371/journal.pcbi.1014571 [code]: SynAPSeg: A novel dataset and image analysis framework for deep learning-based synapse detection and quantification.
Journal: PLoS computational biology
In common: tifffile, TensorFlow, OpenCV, 7 other tools, 3 references
[3] doi:10.1364/boe.605322 [code]: Generalized plaque digitization framework for multi-dimensional mesoscopic images.
Journal: Biomedical optics express
In common: tifffile, TensorFlow, OpenCV, 7 other tools, 2 references
[4] doi:10.1016/j.isci.2026.116206 [code]: Gut distension evokes rapid neural dynamics in vagal and hindbrain populations of larval zebrafish.
Journal: iScience
In common: tifffile, OpenCV, Pillow, 6 other tools, optical imaging (calcium, voltage, 2-photon), 3 references
[5] doi:10.1371/journal.pcbi.1013441 [code]: Large vision model framework for automated C. elegans analysis: From static morphometry to dynamic neural activity.
Journal: PLoS computational biology
In common: tifffile, OpenCV, Pillow, 5 other tools, optical imaging (calcium, voltage, 2-photon), methods / tools, 2 references
[6] doi:10.1016/j.patter.2026.101538 [code]: A multi-modal foundation model for brain disease diagnosis and medical imaging.
Journal: Patterns (New York, N.Y.)
In common: tifffile, TensorFlow, OpenCV, 6 other tools, 2 references
[7] doi:10.1016/j.xpro.2026.104659 [code]: Protocol for simultaneous in vivo two-photon imaging and locomotion quantification during olfactory stimulation and pharmacology in walking Drosophila.
Journal: STAR protocols
In common: tifffile, OpenCV, Pillow, 5 other tools, optical imaging (calcium, voltage, 2-photon), 2 references
[8] doi:10.1016/j.isci.2026.117010 [code]: Deep learning-assisted mapping of dendritic spines using sequential 2D two-photon calcium imaging.
Journal: iScience
In common: tifffile, TensorFlow, OpenCV, 6 other tools, optical imaging (calcium, voltage, 2-photon), 1 reference
[9] doi:10.7554/elife.111876 [code]: Distinct sensorimotor encoding in tuft dendrites and somata associated with action, correction, and learning.
Journal: eLife
In common: tifffile, Pillow, PyTorch, 5 other tools, 3 references
[10] doi:10.1126/sciadv.adv3770 [code]: Evolution of a central dopamine circuit underlies adaptation of a light-evoked sensorimotor response in the blind cavefish.
Journal: Science advances
In common: tifffile, TensorFlow, OpenCV, 6 other tools, 1 reference

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.

NeuroSeg-MF: robust neuron segmentation in two-photon Ca<sup>2+</sup> imaging using multi-feature fusion and detection-guided SAM.

Code ↔ Paper

Paper

The authors' code

Overview

Abstract

Repository

XZH-James/NeuroSeg-MF

Tracing map

Data

Code and data availability statement

Versions

Version 1, 27 September 2026: the first record

Cite

This paper

Similar papers

Contribute

Claim this paper

Correct its record

Validate its tracing map

Add the badge to its README

Request its removal

Discussion, reproductions, activity

NeuroSeg-MF: robust neuron segmentation in two-photon Ca&lt;sup&gt;2+&lt;/sup&gt; imaging using multi-feature fusion and detection-guided SAM.

Code ↔ Paper

Overview

Abstract

Repository

XZH-James/NeuroSeg-MF

Tracing map

Data

Code and data availability statement

Versions

Version 1, 27 September 2026: the first record

Cite

This paper

Similar papers

Contribute

Claim this paper

Correct its record

Validate its tracing map

Add the badge to its README

Request its removal

Discussion, reproductions, activity

NeuroSeg-MF: robust neuron segmentation in two-photon Ca<sup>2+</sup> imaging using multi-feature fusion and detection-guided SAM.