Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import Tensor, nn | |
| import torch.nn.functional as F | |
| import torchvision | |
| from torchvision import transforms | |
| from PIL import Image | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from sklearn.decomposition import PCA | |
| class RandomAffineAndRetMat(torch.nn.Module): | |
| def __init__( | |
| self, | |
| degrees, | |
| translate=None, | |
| scale=None, | |
| shear=None, | |
| interpolation=torchvision.transforms.InterpolationMode.NEAREST, | |
| fill=0, | |
| center=None, | |
| ): | |
| super().__init__() | |
| self.degrees = degrees | |
| self.translate = translate | |
| self.scale = scale | |
| self.shear = shear | |
| self.interpolation = interpolation | |
| self.fill = fill | |
| self.center = center | |
| def forward(self, img): | |
| """ | |
| img (PIL Image or Tensor): Image to be transformed. | |
| Returns: | |
| PIL Image or Tensor: Affine transformed image. | |
| """ | |
| fill = self.fill | |
| if isinstance(img, Tensor): | |
| if isinstance(fill, (int, float)): | |
| fill = [float(fill)] * transforms.functional.get_image_num_channels(img) | |
| else: | |
| fill = [float(f) for f in fill] | |
| img_size = transforms.functional.get_image_size(img) | |
| ret = transforms.RandomAffine.get_params(self.degrees, self.translate, self.scale, self.shear, img_size) | |
| transformed_image = transforms.functional.affine(img, *ret, interpolation=self.interpolation, fill=fill, center=self.center) | |
| affine_matrix = self.get_affine_matrix_from_params(ret) | |
| return transformed_image, affine_matrix | |
| def get_affine_matrix_from_params(self, params): | |
| degrees, translate, scale, shear = params | |
| degrees = torch.tensor(degrees) | |
| shear = torch.tensor(shear) | |
| # パラメータを変換行列に変換 | |
| rotation_matrix = torch.tensor([[torch.cos(torch.deg2rad(degrees)), -torch.sin(torch.deg2rad(degrees)), 0], | |
| [torch.sin(torch.deg2rad(degrees)), torch.cos(torch.deg2rad(degrees)), 0], | |
| [0, 0, 1]]) | |
| translation_matrix = torch.tensor([[1, 0, translate[0]], | |
| [0, 1, translate[1]], | |
| [0, 0, 1]]).to(torch.float32) | |
| scaling_matrix = torch.tensor([[scale, 0, 0], | |
| [0, scale, 0], | |
| [0, 0, 1]]) | |
| shearing_matrix = torch.tensor([[1, -torch.tan(torch.deg2rad(shear[0])), 0], | |
| [-torch.tan(torch.deg2rad(shear[1])), 1, 0], | |
| [0, 0, 1]]) | |
| # 変換行列を合成 | |
| affine_matrix = translation_matrix.mm(rotation_matrix).mm(scaling_matrix).mm(shearing_matrix) | |
| return affine_matrix | |