Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,709 Bytes
02ba63a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
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
|