import argparse import torch from torch import nn from torch.nn import Parameter import torch.nn.functional as F import math from torch import Tensor import numpy as np from typing import Tuple, Union, List, Optional from torch.nn.modules.utils import _pair import torchvision from PIL import Image import json import tqdm from vis import plot_bounding_boxes def letterbox( img: Image.Image, new_shape=(640, 640), color=(114, 114, 114), auto=False, scale_fill=False, scale_up=True ): """ 将 PIL 图像进行 letterbox 处理，用于 YOLOv5 推理前的预处理 来源：YOLOv5 utils/augmentations.py 参数: img (PIL.Image): 输入图像 new_shape (tuple): 模型输入尺寸 (height, width) color (tuple): 填充颜色 (R, G, B) auto (bool): 是否自动选择最小步长（如 YOLOv5 的 auto-anchor） scale_fill (bool): 是否拉伸填充（不保持比例） scale_up (bool): 是否放大图像 返回: image (PIL.Image): letterbox 后的图像 ratio (float): 缩放比例 (dw, dh) (float, float): 左/右 和 上/下 填充的一半（用于还原 bbox） """ shape = img.size # (w, h) new_shape = (new_shape[1], new_shape[0]) # YOLOv5 使用 (w, h) # 计算缩放比例 r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scale_up: # 只缩小，不放大 r = min(r, 1.0) # 计算缩放后的新尺寸 new_unpad = (int(round(shape[0] * r)), int(round(shape[1] * r))) # 缩放图像 img_resized = img.resize(new_unpad, Image.Resampling.BILINEAR) # 填充颜色 top, left = 0, 0 # if auto: # YOLOv5 自动模式：保证尺寸是 32 的倍数 # dw, dh = new_shape[0] - new_unpad[0], new_shape[1] - new_unpad[1] # dw %= 32 # dh %= 32 # elif scale_fill: # # 拉伸填充（不推荐，破坏宽高比） # img_resized = img.resize(new_shape, Image.Resampling.BILINEAR) # dw, dh = 0, 0 # new_unpad = new_shape # else: # 正常 letterbox：居中填充 dw, dh = new_shape[0] - new_unpad[0], new_shape[1] - new_unpad[1] left = dw // 2 top = dh // 2 # 创建新图像并填充背景色 img_letterboxed = Image.new("RGB", new_shape, color) img_letterboxed.paste(img_resized, (left, top)) # 计算填充量（用于还原检测框） ratio = r dw /= 2 # 左侧填充 dh /= 2 # 上侧填充 return img_letterboxed, ratio, (dw, dh) def filter_scores_and_topk(scores, score_thr, topk, results=None): """Filter results using score threshold and topk candidates. Args: scores (Tensor): The scores, shape (num_bboxes, K). score_thr (float): The score filter threshold. topk (int): The number of topk candidates. results (dict or list or Tensor, Optional): The results to which the filtering rule is to be applied. The shape of each item is (num_bboxes, N). Returns: tuple: Filtered results - scores (Tensor): The scores after being filtered, \ shape (num_bboxes_filtered, ). - labels (Tensor): The class labels, shape \ (num_bboxes_filtered, ). - anchor_idxs (Tensor): The anchor indexes, shape \ (num_bboxes_filtered, ). - filtered_results (dict or list or Tensor, Optional): \ The filtered results. The shape of each item is \ (num_bboxes_filtered, N). """ valid_mask = scores > score_thr scores = scores[valid_mask] valid_idxs = torch.nonzero(valid_mask) num_topk = min(topk, valid_idxs.size(0)) # torch.sort is actually faster than .topk (at least on GPUs) scores, idxs = scores.sort(descending=True) scores = scores[:num_topk] topk_idxs = valid_idxs[idxs[:num_topk]] keep_idxs, labels = topk_idxs.unbind(dim=1) filtered_results = None if results is not None: if isinstance(results, dict): filtered_results = {k: v[keep_idxs] for k, v in results.items()} elif isinstance(results, list): filtered_results = [result[keep_idxs] for result in results] elif isinstance(results, torch.Tensor): filtered_results = results[keep_idxs] else: raise NotImplementedError(f'Only supports dict or list or Tensor, ' f'but get {type(results)}.') return scores, labels, keep_idxs, filtered_results ################################################# # ConvNext Backbone ################################################# class Block(nn.Module): r"""ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back We use (2) as we find it slightly faster in PyTorch Args: dim (int): Number of input channels. drop_path (float): Stochastic depth rate. Default: 0.0 layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. """ def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6): super().__init__() self.dwconv = nn.Conv2d( dim, dim, kernel_size=7, padding=3, groups=dim ) # depthwise conv self.norm = LayerNorm(dim, eps=1e-6) self.pwconv1 = nn.Linear( dim, 4 * dim ) # pointwise/1x1 convs, implemented with linear layers self.act = nn.GELU() self.pwconv2 = nn.Linear(4 * dim, dim) self.gamma = ( nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) if layer_scale_init_value > 0 else None ) def forward(self, x: torch.Tensor): input = x x = self.dwconv(x) x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) x = self.norm(x) x = self.pwconv1(x) x = self.act(x) x = self.pwconv2(x) if self.gamma is not None: x = self.gamma * x x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) x = input + x return x class LayerNorm(nn.Module): r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). """ def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): super().__init__() self.weight = nn.Parameter(torch.ones(normalized_shape)) self.bias = nn.Parameter(torch.zeros(normalized_shape)) self.eps = eps self.data_format = data_format if self.data_format not in ["channels_last", "channels_first"]: raise NotImplementedError self.normalized_shape = (normalized_shape,) def forward(self, x: torch.Tensor): if self.data_format == "channels_last": return F.layer_norm( x, self.normalized_shape, self.weight, self.bias, self.eps ) elif self.data_format == "channels_first": u = x.mean(1, keepdim=True) s = (x - u).pow(2).mean(1, keepdim=True) x = (x - u) / torch.sqrt(s + self.eps) x = self.weight[:, None, None] * x + self.bias[:, None, None] return x class ConvNeXt(nn.Module): r"""ConvNeXt A PyTorch impl of : `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf Args: in_chans (int): Number of input image channels. Default: 3 num_classes (int): Number of classes for classification head. Default: 1000 depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] drop_path_rate (float): Stochastic depth rate. Default: 0. layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. """ def __init__( self, model_name ): super().__init__() if model_name == "base": depths = [3, 3, 27, 3] dims = [128, 256, 512, 1024] if model_name == "large": depths = [3, 3, 27, 3] dims = [192, 384, 768, 1536] if model_name == "small": depths = [3, 3, 27, 3] dims = [96, 192, 384, 768] self.downsample_layers = ( nn.ModuleList() ) # stem and 3 intermediate downsampling conv layers stem = nn.Sequential( nn.Conv2d(3, dims[0], kernel_size=4, stride=4), LayerNorm(dims[0], eps=1e-6, data_format="channels_first"), ) self.downsample_layers.append(stem) for i in range(3): downsample_layer = nn.Sequential( LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), ) self.downsample_layers.append(downsample_layer) self.stages = ( nn.ModuleList() ) # 4 feature resolution stages, each consisting of multiple residual blocks dp_rates = [x.item() for x in torch.linspace(0, 0.0, sum(depths))] cur = 0 for i in range(4): stage = nn.Sequential( *[ Block( dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=1e-6, ) for j in range(depths[i]) ] ) self.stages.append(stage) cur += depths[i] def forward(self, x): outputs = [] c1 = self.downsample_layers[0](x) c1 = self.stages[0](c1) outputs.append(c1) c2 = self.downsample_layers[1](c1) c2 = self.stages[1](c2) outputs.append(c2) c3 = self.downsample_layers[2](c2) c3 = self.stages[2](c3) outputs.append(c3) c4 = self.downsample_layers[3](c3) c4 = self.stages[3](c4) outputs.append(c4) return (c1, c2, c3, c4) ################################################# # neck ################################################# activation_table = {'relu':nn.ReLU(), 'silu':nn.SiLU(), 'hardswish':nn.Hardswish() } class ConvModule_torch(nn.Module): '''A combination of Conv + BN + Activation''' def __init__(self, in_channels, out_channels, kernel_size, stride, activation_type, padding=None, groups=1, bias=False): super().__init__() if padding is None: padding = kernel_size // 2 self.conv = nn.Conv2d( in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=bias, ) self.bn = nn.BatchNorm2d(out_channels) if activation_type is not None: self.act = activation_table.get(activation_type) self.activation_type = activation_type def forward(self, x): if self.activation_type is None: return self.bn(self.conv(x)) return self.act(self.bn(self.conv(x))) def forward_fuse(self, x): if self.activation_type is None: return self.conv(x) return self.act(self.conv(x)) class ConvBNReLU(nn.Module): '''Conv and BN with ReLU activation''' def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=None, groups=1, bias=False): super().__init__() self.block = ConvModule_torch(in_channels, out_channels, kernel_size, stride, 'relu', padding, groups, bias) def forward(self, x): return self.block(x) class ConvBNSiLU(nn.Module): '''Conv and BN with SiLU activation''' def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=None, groups=1, bias=False): super().__init__() self.block = ConvModule_torch(in_channels, out_channels, kernel_size, stride, 'silu', padding, groups, bias) def forward(self, x): return self.block(x) class RepBlock(nn.Module): ''' RepBlock is a stage block with rep-style basic block ''' def __init__(self, in_channels, out_channels, block, basic_block, n=1): super().__init__() self.conv1 = BottleRep(in_channels, out_channels, basic_block=basic_block, weight=True) n = n // 2 self.block = nn.Sequential(*(BottleRep(out_channels, out_channels, basic_block=basic_block, weight=True) for _ in range(n - 1))) if n > 1 else None def forward(self, x): x = self.conv1(x) if self.block is not None: x = self.block(x) return x class BottleRep(nn.Module): def __init__(self, in_channels, out_channels, basic_block, weight=False): super().__init__() self.conv1 = basic_block(in_channels, out_channels) self.conv2 = basic_block(out_channels, out_channels) if in_channels != out_channels: self.shortcut = False else: self.shortcut = True if weight: self.alpha = Parameter(torch.ones(1)) else: self.alpha = 1.0 def forward(self, x): outputs = self.conv1(x) outputs = self.conv2(outputs) return outputs + self.alpha * x if self.shortcut else outputs class BepC3(nn.Module): '''CSPStackRep Block''' def __init__(self, in_channels, out_channels, n=1, e=0.5): super().__init__() c_ = int(out_channels * e) # hidden channels self.cv1 = ConvBNReLU(in_channels, c_, 1, 1) self.cv2 = ConvBNReLU(in_channels, c_, 1, 1) self.cv3 = ConvBNReLU(2 * c_, out_channels, 1, 1) self.cv1 = ConvBNSiLU(in_channels, c_, 1, 1) self.cv2 = ConvBNSiLU(in_channels, c_, 1, 1) self.cv3 = ConvBNSiLU(2 * c_, out_channels, 1, 1) self.m = RepBlock(in_channels=c_, out_channels=c_, n=n, block=BottleRep, basic_block=ConvBNSiLU) def forward(self, x): return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1)) class Transpose(nn.Module): '''Normal Transpose, default for upsampling''' def __init__(self, in_channels, out_channels, kernel_size=2, stride=2): super().__init__() self.upsample_transpose = torch.nn.ConvTranspose2d( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, bias=True ) def forward(self, x): return self.upsample_transpose(x) class BiFusion(nn.Module): '''BiFusion Block in PAN''' def __init__(self, in_channels, out_channels): super().__init__() self.cv1 = ConvBNReLU(in_channels[0], out_channels, 1, 1) self.cv2 = ConvBNReLU(in_channels[1], out_channels, 1, 1) self.cv3 = ConvBNReLU(out_channels * 3, out_channels, 1, 1) self.upsample = Transpose( in_channels=out_channels, out_channels=out_channels, ) self.downsample = ConvBNReLU( in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=2 ) def forward(self, x): x0 = self.upsample(x[0]) x1 = self.cv1(x[1]) x2 = self.downsample(self.cv2(x[2])) return self.cv3(torch.cat((x0, x1, x2), dim=1)) class CSPRepBiFPANNeck(nn.Module): """ CSPRepBiFPANNeck module. """ def __init__(self, scale_factor): super().__init__() channels_list=[64, 128, 256, 512, 1024, 256, 128, 128, 256, 256, 512] num_repeats=[1, 6, 12, 18, 6, 12, 12, 12, 12] csp_e=float(1)/2 assert channels_list is not None assert num_repeats is not None stage_block = BepC3 self.reduce_layer0 = ConvBNReLU( in_channels=int(channels_list[4] * scale_factor), # 1024 out_channels=int(channels_list[5] * scale_factor), # 256 kernel_size=1, stride=1 ) self.Bifusion0 = BiFusion( in_channels=[int(channels_list[3] * scale_factor), int(channels_list[2] * scale_factor)], # 512, 256 out_channels=int(channels_list[5] * scale_factor), # 256 ) self.Rep_p4 = stage_block( in_channels=int(channels_list[5] * scale_factor), # 256 out_channels=int(channels_list[5] * scale_factor), # 256 n=num_repeats[5], e=csp_e, ) self.reduce_layer1 = ConvBNReLU( in_channels=int(channels_list[5] * scale_factor), # 256 out_channels=int(channels_list[6] * scale_factor), # 128 kernel_size=1, stride=1 ) self.Bifusion1 = BiFusion( in_channels=[int(channels_list[2] * scale_factor), int(channels_list[1] * scale_factor)], # 256, 128 out_channels=int(channels_list[6] * scale_factor), # 128 ) self.Rep_p3 = stage_block( in_channels=int(channels_list[6] * scale_factor), # 128 out_channels=int(channels_list[6] * scale_factor), # 128 n=num_repeats[6], e=csp_e, ) self.downsample2 = ConvBNReLU( in_channels=int(channels_list[6] * scale_factor), # 128 out_channels=int(channels_list[7] * scale_factor), # 128 kernel_size=3, stride=2 ) self.Rep_n3 = stage_block( in_channels=int(channels_list[6] * scale_factor) + int(channels_list[7] * scale_factor), # 128 + 128 out_channels=int(channels_list[8] * scale_factor), # 256 n=num_repeats[7], e=csp_e, ) self.downsample1 = ConvBNReLU( in_channels=int(channels_list[8] * scale_factor), # 256 out_channels=int(channels_list[9] * scale_factor), # 256 kernel_size=3, stride=2 ) self.Rep_n4 = stage_block( in_channels=int(channels_list[5] * scale_factor) + int(channels_list[9] * scale_factor), # 256 + 256 out_channels=int(channels_list[10] * scale_factor), # 512 n=num_repeats[8], e=csp_e, ) def forward(self, input): (x3, x2, x1, x0) = input fpn_out0 = self.reduce_layer0(x0) f_concat_layer0 = self.Bifusion0([fpn_out0, x1, x2]) f_out0 = self.Rep_p4(f_concat_layer0) fpn_out1 = self.reduce_layer1(f_out0) f_concat_layer1 = self.Bifusion1([fpn_out1, x2, x3]) pan_out2 = self.Rep_p3(f_concat_layer1) down_feat1 = self.downsample2(pan_out2) p_concat_layer1 = torch.cat([down_feat1, fpn_out1], 1) pan_out1 = self.Rep_n3(p_concat_layer1) down_feat0 = self.downsample1(pan_out1) p_concat_layer2 = torch.cat([down_feat0, fpn_out0], 1) pan_out0 = self.Rep_n4(p_concat_layer2) outputs = [pan_out2, pan_out1, pan_out0] return outputs ################################################# # head ################################################# class BNContrastiveHead(nn.Module): """ Batch Norm Contrastive Head for YOLO-World using batch norm instead of l2-normalization Args: embed_dims (int): embed dim of text and image features norm_cfg (dict): normalization params """ def __init__(self, embed_dims: int, use_einsum: bool = True) -> None: super().__init__() self.norm = nn.BatchNorm2d(embed_dims, momentum=0.03, eps=0.001) self.bias = nn.Parameter(torch.zeros([])) # use -1.0 is more stable self.logit_scale = nn.Parameter(-1.0 * torch.ones([])) self.use_einsum = use_einsum def forward(self, x: Tensor, w: Tensor) -> Tensor: """Forward function of contrastive learning.""" x = self.norm(x) w = F.normalize(w, dim=-1, p=2) if self.use_einsum: x = torch.einsum('bchw,bkc->bkhw', x, w) else: batch, channel, height, width = x.shape _, k, _ = w.shape x = x.permute(0, 2, 3, 1) # bchw->bhwc x = x.reshape(batch, -1, channel) # bhwc->b(hw)c w = w.permute(0, 2, 1) # bkc->bck x = torch.matmul(x, w) x = x.reshape(batch, height, width, k) x = x.permute(0, 3, 1, 2) x = x * self.logit_scale.exp() + self.bias return x class YOLOWorldHeadModule(nn.Module): """Head Module for YOLO-World Args: embed_dims (int): embed dim for text feautures and image features use_bn_head (bool): use batch normalization head """ def __init__(self, embed_dims: int, in_channels, use_bn_head: bool = False, use_einsum: bool = True, freeze_all: bool = False, ) -> None: self.embed_dims = embed_dims self.use_bn_head = use_bn_head self.use_einsum = use_einsum self.freeze_all = freeze_all self.in_channels = in_channels self.reg_max = 16 super().__init__() self._init_layers() def _init_layers(self) -> None: """initialize conv layers in YOLOv8 head.""" # Init decouple head self.cls_preds = nn.ModuleList() self.reg_preds = nn.ModuleList() self.cls_contrasts = nn.ModuleList() cls_out_channels = 256 self.featmap_strides = [8, 16, 32] self.num_levels = len(self.in_channels) reg_out_channels = max( (16, self.in_channels[0] // 4, self.reg_max * 4)) for i in range(self.num_levels): self.reg_preds.append( nn.Sequential( nn.Conv2d(in_channels=self.in_channels[i], out_channels=reg_out_channels, kernel_size=3, stride=1, padding=1, bias=False,), nn.BatchNorm2d(reg_out_channels, momentum=0.03, eps=0.001), nn.SiLU(), nn.Conv2d(in_channels=reg_out_channels, out_channels=reg_out_channels, kernel_size=3, stride=1, padding=1, bias=False,), nn.BatchNorm2d(reg_out_channels, momentum=0.03, eps=0.001), nn.SiLU(), nn.Conv2d(in_channels=reg_out_channels, out_channels=4 * self.reg_max, kernel_size=1))) self.cls_preds.append( nn.Sequential( nn.Conv2d(in_channels=self.in_channels[i], out_channels=cls_out_channels, kernel_size=3, stride=1, padding=1, bias=False,), nn.BatchNorm2d(cls_out_channels, momentum=0.03, eps=0.001), nn.SiLU(), nn.Conv2d(in_channels=cls_out_channels, out_channels=cls_out_channels, kernel_size=3, stride=1, padding=1, bias=False,), nn.BatchNorm2d(cls_out_channels, momentum=0.03, eps=0.001), nn.SiLU(), nn.Conv2d(in_channels=cls_out_channels, out_channels=self.embed_dims, kernel_size=1))) self.cls_contrasts.append(BNContrastiveHead(self.embed_dims, use_einsum=self.use_einsum)) proj = torch.arange(self.reg_max, dtype=torch.float) self.register_buffer('proj', proj, persistent=False) def forward(self, img_feats: Tuple[Tensor], txt_feats: Tensor) -> Tuple[List]: """Forward features from the upstream network.""" assert len(img_feats) == self.num_levels txt_feats = [txt_feats for _ in range(self.num_levels)] outputs = [] for i in range(self.num_levels): outputs.append( self.forward_single(img_feats[i], txt_feats[i], self.cls_preds[i], self.reg_preds[i], self.cls_contrasts[i])) return tuple(outputs) def forward_single(self, img_feat: Tensor, txt_feat: Tensor, cls_pred: nn.ModuleList, reg_pred: nn.ModuleList, cls_contrast: nn.ModuleList) -> Tuple: """Forward feature of a single scale level.""" b, _, h, w = img_feat.shape cls_embed = cls_pred(img_feat) cls_logit = cls_contrast(cls_embed, txt_feat) bbox_dist_preds = reg_pred(img_feat) if self.reg_max > 1: bbox_dist_preds = bbox_dist_preds.reshape( [-1, 4, self.reg_max, h * w]).permute(0, 3, 1, 2) # TODO: The get_flops script cannot handle the situation of # matmul, and needs to be fixed later # bbox_preds = bbox_dist_preds.softmax(3).matmul(self.proj) bbox_preds = bbox_dist_preds.softmax(3).matmul( self.proj.view([-1, 1])).squeeze(-1) bbox_preds = bbox_preds.transpose(1, 2).reshape(b, -1, h, w) else: bbox_preds = bbox_dist_preds if self.training: return cls_logit, bbox_preds, bbox_dist_preds else: return cls_logit, bbox_preds ################################################# # prior generator ################################################# class MlvlPointGenerator: """Standard points generator for multi-level (Mlvl) feature maps in 2D points-based detectors. Args: strides (list[int] | list[tuple[int, int]]): Strides of anchors in multiple feature levels in order (w, h). offset (float): The offset of points, the value is normalized with corresponding stride. Defaults to 0.5. """ def __init__(self, strides: Union[List[int], List[Tuple[int, int]]], offset: float = 0.5) -> None: self.strides = [_pair(stride) for stride in strides] self.offset = offset @property def num_levels(self) -> int: """int: number of feature levels that the generator will be applied""" return len(self.strides) @property def num_base_priors(self) -> List[int]: """list[int]: The number of priors (points) at a point on the feature grid""" return [1 for _ in range(len(self.strides))] def _meshgrid(self, x: Tensor, y: Tensor, row_major: bool = True) -> Tuple[Tensor, Tensor]: yy, xx = torch.meshgrid(y, x) if row_major: # warning .flatten() would cause error in ONNX exporting # have to use reshape here return xx.reshape(-1), yy.reshape(-1) else: return yy.reshape(-1), xx.reshape(-1) def grid_priors(self, featmap_sizes: List[Tuple], dtype: torch.dtype = torch.float32, device = 'cuda', with_stride: bool = False) -> List[Tensor]: """Generate grid points of multiple feature levels. Args: featmap_sizes (list[tuple]): List of feature map sizes in multiple feature levels, each size arrange as as (h, w). dtype (:obj:`dtype`): Dtype of priors. Defaults to torch.float32. device (str | torch.device): The device where the anchors will be put on. with_stride (bool): Whether to concatenate the stride to the last dimension of points. Return: list[torch.Tensor]: Points of multiple feature levels. The sizes of each tensor should be (N, 2) when with stride is ``False``, where N = width * height, width and height are the sizes of the corresponding feature level, and the last dimension 2 represent (coord_x, coord_y), otherwise the shape should be (N, 4), and the last dimension 4 represent (coord_x, coord_y, stride_w, stride_h). """ assert self.num_levels == len(featmap_sizes) multi_level_priors = [] for i in range(self.num_levels): priors = self.single_level_grid_priors( featmap_sizes[i], level_idx=i, dtype=dtype, device=device, with_stride=with_stride) multi_level_priors.append(priors) return multi_level_priors def single_level_grid_priors(self, featmap_size: Tuple[int], level_idx: int, dtype: torch.dtype = torch.float32, device = 'cuda', with_stride: bool = False) -> Tensor: """Generate grid Points of a single level. Note: This function is usually called by method ``self.grid_priors``. Args: featmap_size (tuple[int]): Size of the feature maps, arrange as (h, w). level_idx (int): The index of corresponding feature map level. dtype (:obj:`dtype`): Dtype of priors. Defaults to torch.float32. device (str | torch.device): The device the tensor will be put on. Defaults to 'cuda'. with_stride (bool): Concatenate the stride to the last dimension of points. Return: Tensor: Points of single feature levels. The shape of tensor should be (N, 2) when with stride is ``False``, where N = width * height, width and height are the sizes of the corresponding feature level, and the last dimension 2 represent (coord_x, coord_y), otherwise the shape should be (N, 4), and the last dimension 4 represent (coord_x, coord_y, stride_w, stride_h). """ feat_h, feat_w = featmap_size stride_w, stride_h = self.strides[level_idx] shift_x = (torch.arange(0, feat_w, device=device) + self.offset) * stride_w # keep featmap_size as Tensor instead of int, so that we # can convert to ONNX correctly shift_x = shift_x.to(dtype) shift_y = (torch.arange(0, feat_h, device=device) + self.offset) * stride_h # keep featmap_size as Tensor instead of int, so that we # can convert to ONNX correctly shift_y = shift_y.to(dtype) shift_xx, shift_yy = self._meshgrid(shift_x, shift_y) if not with_stride: shifts = torch.stack([shift_xx, shift_yy], dim=-1) else: # use `shape[0]` instead of `len(shift_xx)` for ONNX export stride_w = shift_xx.new_full((shift_xx.shape[0], ), stride_w).to(dtype) stride_h = shift_xx.new_full((shift_yy.shape[0], ), stride_h).to(dtype) shifts = torch.stack([shift_xx, shift_yy, stride_w, stride_h], dim=-1) all_points = shifts.to(device) return all_points def valid_flags(self, featmap_sizes: List[Tuple[int, int]], pad_shape: Tuple[int], device = 'cuda') -> List[Tensor]: """Generate valid flags of points of multiple feature levels. Args: featmap_sizes (list(tuple)): List of feature map sizes in multiple feature levels, each size arrange as as (h, w). pad_shape (tuple(int)): The padded shape of the image, arrange as (h, w). device (str | torch.device): The device where the anchors will be put on. Return: list(torch.Tensor): Valid flags of points of multiple levels. """ assert self.num_levels == len(featmap_sizes) multi_level_flags = [] for i in range(self.num_levels): point_stride = self.strides[i] feat_h, feat_w = featmap_sizes[i] h, w = pad_shape[:2] valid_feat_h = min(int(np.ceil(h / point_stride[1])), feat_h) valid_feat_w = min(int(np.ceil(w / point_stride[0])), feat_w) flags = self.single_level_valid_flags((feat_h, feat_w), (valid_feat_h, valid_feat_w), device=device) multi_level_flags.append(flags) return multi_level_flags def single_level_valid_flags(self, featmap_size: Tuple[int, int], valid_size: Tuple[int, int], device = 'cuda') -> Tensor: """Generate the valid flags of points of a single feature map. Args: featmap_size (tuple[int]): The size of feature maps, arrange as as (h, w). valid_size (tuple[int]): The valid size of the feature maps. The size arrange as as (h, w). device (str | torch.device): The device where the flags will be put on. Defaults to 'cuda'. Returns: torch.Tensor: The valid flags of each points in a single level \ feature map. """ feat_h, feat_w = featmap_size valid_h, valid_w = valid_size assert valid_h <= feat_h and valid_w <= feat_w valid_x = torch.zeros(feat_w, dtype=torch.bool, device=device) valid_y = torch.zeros(feat_h, dtype=torch.bool, device=device) valid_x[:valid_w] = 1 valid_y[:valid_h] = 1 valid_xx, valid_yy = self._meshgrid(valid_x, valid_y) valid = valid_xx & valid_yy return valid def sparse_priors(self, prior_idxs: Tensor, featmap_size: Tuple[int], level_idx: int, dtype: torch.dtype = torch.float32, device = 'cuda') -> Tensor: """Generate sparse points according to the ``prior_idxs``. Args: prior_idxs (Tensor): The index of corresponding anchors in the feature map. featmap_size (tuple[int]): feature map size arrange as (w, h). level_idx (int): The level index of corresponding feature map. dtype (obj:`torch.dtype`): Date type of points. Defaults to ``torch.float32``. device (str | torch.device): The device where the points is located. Returns: Tensor: Anchor with shape (N, 2), N should be equal to the length of ``prior_idxs``. And last dimension 2 represent (coord_x, coord_y). """ height, width = featmap_size x = (prior_idxs % width + self.offset) * self.strides[level_idx][0] y = ((prior_idxs // width) % height + self.offset) * self.strides[level_idx][1] prioris = torch.stack([x, y], 1).to(dtype) prioris = prioris.to(device) return prioris def distance2bbox( points: Tensor, distance: Tensor, max_shape = None ) -> Tensor: """Decode distance prediction to bounding box. Args: points (Tensor): Shape (B, N, 2) or (N, 2). distance (Tensor): Distance from the given point to 4 boundaries (left, top, right, bottom). Shape (B, N, 4) or (N, 4) max_shape (Union[Sequence[int], Tensor, Sequence[Sequence[int]]], optional): Maximum bounds for boxes, specifies (H, W, C) or (H, W). If priors shape is (B, N, 4), then the max_shape should be a Sequence[Sequence[int]] and the length of max_shape should also be B. Returns: Tensor: Boxes with shape (N, 4) or (B, N, 4) """ x1 = points[..., 0] - distance[..., 0] y1 = points[..., 1] - distance[..., 1] x2 = points[..., 0] + distance[..., 2] y2 = points[..., 1] + distance[..., 3] bboxes = torch.stack([x1, y1, x2, y2], -1) if max_shape is not None: if bboxes.dim() == 2 and not torch.onnx.is_in_onnx_export(): # speed up bboxes[:, 0::2].clamp_(min=0, max=max_shape[1]) bboxes[:, 1::2].clamp_(min=0, max=max_shape[0]) return bboxes if not isinstance(max_shape, torch.Tensor): max_shape = x1.new_tensor(max_shape) max_shape = max_shape[..., :2].type_as(x1) if max_shape.ndim == 2: assert bboxes.ndim == 3 assert max_shape.size(0) == bboxes.size(0) min_xy = x1.new_tensor(0) max_xy = torch.cat([max_shape, max_shape], dim=-1).flip(-1).unsqueeze(-2) bboxes = torch.where(bboxes < min_xy, min_xy, bboxes) bboxes = torch.where(bboxes > max_xy, max_xy, bboxes) return bboxes class SimpleYOLOWorldDetector(nn.Module): """Implementation of YOLO World Series""" def __init__(self, backbone_size, prompt_dim=768, num_prompts=512, num_proposals=300) -> None: super().__init__() self.backbone = ConvNeXt(backbone_size) if backbone_size == 'base': scale_factor = 1.0 in_channels = [128, 256, 512] self.img_size = (640, 640) self.grid_size = [6400, 1600, 400] elif backbone_size == 'large': scale_factor = 1.5 in_channels = [192, 384, 768] self.img_size = (1280, 1280) self.grid_size = [6400*4, 1600*4, 400*4] self.num_proposals = num_proposals self.neck = CSPRepBiFPANNeck(scale_factor) self.bbox_head = YOLOWorldHeadModule(embed_dims=prompt_dim, in_channels=in_channels, use_bn_head=True, use_einsum=True) embeddings = nn.functional.normalize(torch.randn( (num_prompts, prompt_dim)), dim=-1) self.embeddings = nn.Parameter(embeddings) self.prior_generator = MlvlPointGenerator(strides=[8, 16, 32], offset=0.5) def forward(self, image_paths: List[str], rescale=True): inputs = [] ratios = [] offsets = [] ori_shapes = [] for image_path in image_paths: if type(image_path) is str: img = Image.open(image_path).convert("RGB") else: img = image_path width, height = img.size ori_shape = (height, width) ori_shapes.append(ori_shape) img, ratio, offset = letterbox(img, self.img_size) img = torch.tensor(np.array(img)).permute(2, 0, 1).to(self.embeddings.data.dtype) # HWC, RGB img = img / 255.0 inputs.append(img) ratios.append(ratio) offsets.append(offset) inputs = torch.stack(inputs, dim=0).to(self.embeddings.device) img_feats = self.backbone(inputs) img_feats = self.neck(img_feats) results = self.head_predict(img_feats) for i in range(len(results)): # print(results[i]['bboxes']) results[i]['bboxes'] -= results[i]['bboxes'].new_tensor([ offsets[i][0], offsets[i][1], offsets[i][0], offsets[i][1] ]) # print(results[i]['bboxes']) if rescale: results[i]['bboxes'] /= ratios[i] results[i]['bboxes'][:, 0::2] = results[i]['bboxes'][:, 0::2].clamp_(0, ori_shapes[i][1]) results[i]['bboxes'][:, 1::2] = results[i]['bboxes'][:, 1::2].clamp_(0, ori_shapes[i][0]) # print(results[i]['bboxes']) return results def head_module_forward_single( self, img_feat: Tensor, cls_pred: nn.ModuleList, reg_pred: nn.ModuleList, cls_contrast: nn.ModuleList, ): module=self.bbox_head b, _, h, w = img_feat.shape cls_embed = cls_pred(img_feat) cls_embed = cls_contrast.norm(cls_embed) cls_logits = torch.einsum('bchw,kc->bkhw', cls_embed, self.embeddings) cls_logits = cls_logits * cls_contrast.logit_scale.exp() + cls_contrast.bias bbox_dist_preds = reg_pred(img_feat) if module.reg_max > 1: bbox_dist_preds = bbox_dist_preds.reshape( [-1, 4, module.reg_max, h * w] ).permute(0, 3, 1, 2) # TODO: The get_flops script cannot handle the situation of # matmul, and needs to be fixed later # bbox_preds = bbox_dist_preds.softmax(3).matmul(self.proj) bbox_preds = ( bbox_dist_preds.softmax(3).matmul(module.proj.view([-1, 1])).squeeze(-1) ) bbox_preds = bbox_preds.transpose(1, 2).reshape(b, -1, h, w) else: bbox_preds = bbox_dist_preds return cls_embed, bbox_preds,cls_logits def head_predict(self, img_feats): bbox_embed, bbox_preds, cls_scores = [], [], [] for i in range(len(img_feats)): box_embed, bbox_pred, cls_score = self.head_module_forward_single( img_feats[i], self.bbox_head.cls_preds[i], self.bbox_head.reg_preds[i], self.bbox_head.cls_contrasts[i], ) bbox_embed.append(box_embed) bbox_preds.append(bbox_pred) cls_scores.append(cls_score) # objectness = None, with_objectness=False, multi_label=False txt_channel = bbox_embed[0].shape[1] num_imgs = bbox_embed[0].shape[0] featmap_sizes = [x.shape[2:] for x in bbox_preds] mlvl_priors = self.prior_generator.grid_priors( featmap_sizes, dtype=bbox_embed[0].dtype, device=bbox_embed[0].device ) flatten_priors = torch.cat(mlvl_priors) mlvl_strides = [ flatten_priors.new_full((featmap_size.numel(),), stride) for featmap_size, stride in zip(featmap_sizes, self.bbox_head.featmap_strides) ] flatten_stride = torch.cat(mlvl_strides) flatten_bbox_embed = [ x.permute(0, 2, 3, 1).reshape(num_imgs, -1, txt_channel) for x in bbox_embed ] flatten_cls_scores = [ cls_score.permute(0, 2, 3, 1).reshape(num_imgs, -1, cls_scores[0].shape[1]) for cls_score in cls_scores ] flatten_cls_scores = torch.cat(flatten_cls_scores, dim=1).sigmoid() flatten_bbox_preds = [ bbox_pred.permute(0, 2, 3, 1).reshape(num_imgs, -1, 4) for bbox_pred in bbox_preds ] flatten_bbox_embed = torch.cat(flatten_bbox_embed, dim=1) flatten_bbox_preds = torch.cat(flatten_bbox_preds, dim=1) flatten_bbox_preds = flatten_bbox_preds * flatten_stride[None, :, None] flatten_decoded_bbox = distance2bbox( flatten_priors[None], flatten_bbox_preds ) results_list = [] for bbox, embed,scores in zip( flatten_decoded_bbox, flatten_bbox_embed,flatten_cls_scores ): original_scores = scores.clone() scores, labels, keep_idxs, _ = filter_scores_and_topk( scores, 0.0, 30000) # if len(scores) < 50: # scores, labels, keep_idxs, _ = filter_scores_and_topk( # original_scores, 0.0001, 50) bbox = bbox[keep_idxs] embed = embed[keep_idxs] idx = torchvision.ops.batched_nms(bbox.float(), scores.float(), labels, 0.7)[:self.num_proposals] bbox = bbox[idx] embed = embed[idx] results_list.append({ 'bboxes': bbox, 'embeddings': embed, 'scores': scores[idx], }) return results_list if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--wedetect_uni_checkpoint', type=str, default='') parser.add_argument('--image', type=str, default='') parser.add_argument('--score_thre', type=float, default=0.1) parser.add_argument('--num_proposals', type=int, default=300) parser.add_argument('--visualize', action='store_true') args = parser.parse_args() # load detection model model_size = 'base' if 'base' in args.wedetect_uni_checkpoint else 'large' model = SimpleYOLOWorldDetector(backbone_size=model_size, prompt_dim=768, num_prompts=256, num_proposals=args.num_proposals) checkpoint = torch.load(args.wedetect_uni_checkpoint, map_location='cpu') # backbone keys = list(checkpoint.keys()) for key in keys: if 'backbone' in key: new_key = key.replace('backbone.image_model.model.', 'backbone.') checkpoint[new_key] = checkpoint.pop(key) # head keys = list(checkpoint.keys()) for key in keys: if 'bbox_head' in key: new_key = key.replace('bbox_head.head_module.', 'bbox_head.') new_key = new_key.replace('0.2.', '0.6.') new_key = new_key.replace('1.2.', '1.6.') new_key = new_key.replace('2.2.', '2.6.') new_key = new_key.replace('1.bn', '4') new_key = new_key.replace('1.conv', '3') new_key = new_key.replace('0.bn', '1') new_key = new_key.replace('0.conv', '0') checkpoint[new_key] = checkpoint.pop(key) msg = model.load_state_dict(checkpoint, strict=False) print(msg) model = model.cuda() model.eval() with torch.no_grad(): outputs = model([args.image]) pred_bboxes = outputs[0]['bboxes'].float().cpu() pred_scores = outputs[0]['scores'].float().cpu() if args.score_thre > 0: mask = pred_scores > args.score_thre pred_bboxes = pred_bboxes[mask] pred_scores = pred_scores[mask] if args.visualize: image = Image.open(args.image) pred_image = plot_bounding_boxes(image, pred_bboxes.cpu().tolist()) pred_image.save("pred.png") # 你可以自定义保存路径和文件名