现在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,这是个复制边缘像素,对于背景简单的就还可以