YoloV8修改分类(Classify)的前处理（记录）

news2024/11/15 2:04:25

修改原因

yolo自带的分类前处理对于长方形的数据不够友好，存在特征丢失等问题
修改后虽然解决了这个问题但是局部特征也会丢失因为会下采样程度多于自带的，总之具体哪种好不同数据应该表现不同
我的数据中大量长宽比很大的数据所以尝试修改自带的前处理，以保证理论上的合理性。

修改过程

yolo中自带的分类前处理和检测有一些差异

调试推理代码发现ultralytics/models/yolo/classify/predict.py中对图像进行前处理的操作主要是self.transforms

def preprocess(self, img):
        """Converts input image to model-compatible data type."""
    if not isinstance(img, torch.Tensor):
        is_legacy_transform = any(
            self._legacy_transform_name in str(transform) for transform in self.transforms.transforms
        )
        if is_legacy_transform:  # to handle legacy transforms
            img = torch.stack([self.transforms(im) for im in img], dim=0)
        else:
            # import ipdb;ipdb.set_trace()
            img = torch.stack(
                [self.transforms(Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))) for im in img], dim=0
            )
    img = (img if isinstance(img, torch.Tensor) else torch.from_numpy(img)).to(self.model.device)
    return img.half() if self.model.fp16 else img.float()  # uint8 to fp16/32

通过调试打印self.transforms得到

Compose(
    Resize(size=96, interpolation=bilinear, max_size=None, antialias=True)
    CenterCrop(size=(96, 96))
    ToTensor()
    Normalize(mean=tensor([0., 0., 0.]), std=tensor([1., 1., 1.]))
)

假设我设置的imgsz为96，从这里简单的解读可以理解为先进行resize然后进行中心裁切保证输入尺寸为96x96

具体的查看哪里可以修改前处理，首先发现在ultralytics/engine/predictor.py中

def setup_source(self, source):
 """Sets up source and inference mode."""
  self.imgsz = check_imgsz(self.args.imgsz, stride=self.model.stride, min_dim=2)  # check image size
  # import ipdb; ipdb.set_trace()
  self.transforms = (
      getattr(
          self.model.model,
          "transforms",
          classify_transforms(self.imgsz[0], crop_fraction=self.args.crop_fraction), #dujiang
      )
      if self.args.task == "classify"
      else None
  )

可以发现self.transforms主要调用的是classify_transforms方法

进一步我们在ultralytics/data/augment.py中找到`classify_transforms`的实现

if scale_size[0] == scale_size[1]:
      # Simple case, use torchvision built-in Resize with the shortest edge mode (scalar size arg)
      tfl = [T.Resize(scale_size[0], interpolation=getattr(T.InterpolationMode, interpolation))]
  else:
      # Resize the shortest edge to matching target dim for non-square target
      tfl = [T.Resize(scale_size)]
  tfl.extend(
      [
          T.CenterCrop(size),
          T.ToTensor(),
          T.Normalize(mean=torch.tensor(mean), std=torch.tensor(std)),
      ]
  )

发现和我们的设想基本一致，查看代码逻辑首先是针对正方形图像会将图像缩放到指定的高度，同时保持长宽比，确保较短的一边正好等于目标尺寸，非正方形图片将短边resize到指定大小，长边此时可能是超出的，所以 T.CenterCrop(size)进行中心裁切确保尺寸是我们指定的

针对上面的分析可能问题就很明显了，如果处理的图像是长宽比非常不均匀的图像，那么中心裁切会导致丢失大量信息，我参考了检测的方法，决定将分类的预处理修改为填充而不是裁切。

首先确定思想，我想做的是根据长边resize到指定尺寸并且保证长宽比，短边会不足，刚好与原本的代码逻辑相反
然后短边不足的地方进行填充保证短边也达到指定尺寸（填充yolo好像一般是144，这里我也选择144）
具体实现如下

添加两个类分别实现resize和padding

class ResizeLongestSide:
    def __init__(self, size, interpolation):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img):
        # 获取图像的当前尺寸
        width, height = img.size
        # 计算缩放比例
        if width > height:
            new_width = self.size
            new_height = int(self.size * height / width)
        else:
            new_height = self.size
            new_width = int(self.size * width / height)
        # 按长边缩放
        return img.resize((new_width, new_height), Image.BILINEAR)

class PadToSquare:
    def __init__(self, size, fill=(114)):
        self.size = size
        self.fill = fill

    def __call__(self, img):
        # 获取当前尺寸
        width, height = img.size
        # 计算需要填充的大小
        delta_w = self.size - width
        delta_h = self.size - height
        padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
        # 填充图像
        return F.pad(img, padding, fill=self.fill, padding_mode='constant')

调用上面的类进行实现

def classify_transforms(
    size=96,
    mean=DEFAULT_MEAN,
    std=DEFAULT_STD,
    interpolation="BILINEAR",
    crop_fraction: float = DEFAULT_CROP_FRACTION,
    padding_color=(114, 114, 114),  # 默认填充为灰色
):

    import torchvision.transforms as T
    import torch
    from torchvision.transforms import functional as F

    # import ipdb;ipdb.set_trace()
    tfl = [
        # T.ClassifyLetterBox(size),
        ResizeLongestSide(size, interpolation=getattr(T.InterpolationMode, interpolation)),  # 按长边缩放
        PadToSquare(size, fill=padding_color),  # 填充至正方形
        T.ToTensor(),
        T.Normalize(mean=torch.tensor(mean), std=torch.tensor(std)),
    ]

    return T.Compose(tfl)

想要训练前先确定自己修改是否符合预期进行如下测试

Examples:
    >>> from ultralytics.data.augment import LetterBox, classify_transforms, classify_transforms_with_padding
    >>> from PIL import Image
    >>> transforms = classify_transforms_with_padding(size=96)
    >>> img = Image.open('bus.jpg')  3ch img_rgb = Image.merge('RGB', (img, img, img))
    >>> transformed_img = transforms(img)
    >>>import torchvision.transforms as T
    >>>DEFAULT_MEAN = (0.0, 0.0, 0.0)
    >>>DEFAULT_STD = (1.0, 1.0, 1.0)
    >>>import torch
    >>>def save_transformed_image(transformed_img, save_path="transformed_image.png"):
    # 定义反向变换，将张量转换回 PIL 图像
    unnormalize = T.Normalize(
        mean=[-m / s for m, s in zip(DEFAULT_MEAN, DEFAULT_STD)],
        std=[1 / s for s in DEFAULT_STD]
    )
    img_tensor = unnormalize(transformed_img)
    img_tensor = torch.clamp(img_tensor, 0, 1)
    to_pil = T.ToPILImage()
    img_pil = to_pil(img_tensor)
    img_pil.save(save_path)
    print(f"Image saved at {save_path}")
    >>>save_transformed_image(transformed_img, save_path="transformed_image.png")