现在transform的包还挺多,不过我没搜到/不会搜他们有能满足我需求的方法,所以还是摸索这个基本的torchvision版本

class RandomRotateWithPadding:
    def __init__(self, degrees, padding=0, fill=0, padding_mode='edge'):
        self.degrees = degrees
        self.padding = padding
        self.padding_mode = padding_mode
        self.fill = fill

    def __call__(self, img):
        # 计算旋转角度
        angle = random.uniform(self.degrees[0], self.degrees[1])

        # 将PIL图像转换为numpy数组
        np_img = np.array(img)

        # 计算需要的padding大小
        diagonal = int((np_img.shape[0] ** 2 + np_img.shape[1] ** 2) ** 0.5)
        padding = (diagonal - min(np_img.shape[0], np_img.shape[1])) // 2

        # 使用numpy的pad函数应用reflect填充
        np_padded = np.pad(np_img, [(padding, padding), (padding, padding), (0, 0)], mode=self.padding_mode)

        # 将填充后的numpy数组转回PIL图像
        img_padded = Image.fromarray(np_padded)

        # 旋转图像
        img_rotated = img_padded.rotate(angle, expand=False)

        # 裁剪回原始尺寸
        x0 = y0 = padding
        x1 = x0 + img.width
        y1 = y0 + img.height
        img_cropped = img_rotated.crop((x0, y0, x1, y1))

        return img_cropped


这里RandomRotateWithPadding(degrees=(-30, 30), padding_mode="edge")还是用edge还一些,可以比较平滑的延展原图的背景,其他像reflect就是单纯镜像,不行


class RandomRotateWithOpenCVPadding:
    def __init__(self, degrees, padding_mode=cv2.BORDER_REFLECT):
        self.degrees = degrees
        self.padding_mode = padding_mode

    def __call__(self, img):
        angle = random.uniform(*self.degrees)
        # 将PIL图像转换为OpenCV图像
        img_cv = np.array(img)
        img_cv = img_cv[:, :, ::-1]  # PIL到OpenCV颜色通道转换:RGB到BGR

        # 计算需要的padding大小
        diagonal = int((img_cv.shape[0]**2 + img_cv.shape[1]**2)**0.5)
        padding = (diagonal - min(img_cv.shape[0], img_cv.shape[1])) // 2

        # 使用OpenCV的copyMakeBorder应用padding
        img_padded = cv2.copyMakeBorder(img_cv, padding, padding, padding, padding, self.padding_mode)

        # 创建旋转矩阵并应用旋转
        center = (diagonal // 2, diagonal // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        img_rotated = cv2.warpAffine(img_padded, M, (diagonal, diagonal))

        # 裁剪回原始尺寸
        x0 = y0 = padding
        x1 = x0 + img.width
        y1 = y0 + img.height
        img_cropped = img_rotated[y0:y1, x0:x1]

        # 将OpenCV图像转换回PIL图像
        img_cropped = cv2.cvtColor(img_cropped, cv2.COLOR_BGR2RGB)
        img_final = Image.fromarray(img_cropped)

        return img_final

捯饬半天干脆试试cv2咋样,这里RandomRotateWithOpenCVPadding(degrees=(-30, 30), padding_mode=cv2.BORDER_REPLICATE),好像没有edge所以最好是用cv2.BORDER_REPLICATE,这是个复制边缘像素,对于背景简单的就还可以

04-30 12:32