utils.py

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

def norm_img(img):
  return (img-img.min())/(img.max()-img.min())

def preprocess_uploaded_image(uploaded_image, image_size):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # ndarrayの場合はPILイメージに変換
    if type(uploaded_image) == np.ndarray:
        uploaded_image = Image.fromarray(uploaded_image)
    uploaded_image = uploaded_image.convert("RGB")
    uploaded_image = uploaded_image.resize((image_size, image_size))
    uploaded_image = np.array(uploaded_image).transpose(2, 0, 1) / 255.0
    uploaded_image = torch.tensor(uploaded_image, dtype=torch.float32).unsqueeze(0).to(device)
    return uploaded_image

def get_heatmaps(img, feature_map, source_num, x_coords, y_coords, uploaded_image):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    image_size = img.size(2)
    
    batch_size = feature_map.size(0)
    feature_dim = feature_map.size(1)
    
    target_num = batch_size - 1

    x_coords = [x_coords] * batch_size
    y_coords = [y_coords] * batch_size

    vectors = feature_map[torch.arange(feature_map.size(0)), :, y_coords, x_coords]
    vector = vectors[source_num]

    reshaped_feature_map = feature_map.permute(0, 2, 3, 1).view(feature_map.size(0), -1, feature_dim)
    batch_distance_map = F.pairwise_distance(reshaped_feature_map, vector).view(feature_map.size(0), image_size, image_size)
    
    norm_batch_distance_map = 1 / torch.cosh(20 * (batch_distance_map - batch_distance_map.min()) / (batch_distance_map.max() - batch_distance_map.min())) ** 2

    source_map = norm_batch_distance_map[source_num].detach().cpu()
    target_map = norm_batch_distance_map[target_num].detach().cpu()

    alpha = 0.7
    blended_source = (1 - alpha) * img[source_num] + alpha * torch.cat(((norm_batch_distance_map[source_num] / norm_batch_distance_map[source_num].max()).unsqueeze(0), torch.zeros(2, image_size, image_size, device=device)))
    blended_target = (1 - alpha) * img[target_num] + alpha * torch.cat(((norm_batch_distance_map[target_num] / norm_batch_distance_map[target_num].max()).unsqueeze(0), torch.zeros(2, image_size, image_size, device=device)))
    
    blended_source = blended_source.detach().cpu()
    blended_target = blended_target.detach().cpu()
    
    return source_map, target_map, blended_source, blended_target

def get_mean_vector(feature_map, points):
  keypoints_size = points.size(1)

  mean_vector_list = []
  for i in range(keypoints_size):
      x_coords, y_coords = torch.round(points[:,i].t()).to(torch.long)
      vectors = feature_map[torch.arange(feature_map.size(0)), :, y_coords, x_coords]  # 1次元ベクトルに合わせてサイズを調整
      # mean_vector = vectors[0:10].mean(0) # 10個の特徴マップの平均ベクトルを取得
      mean_vector = vectors.mean(0).detach().cpu().numpy()
      mean_vector_list.append(mean_vector)
  return mean_vector_list

def get_keypoint_heatmaps(feature_map, mean_vector_list, keypoints_size, imgs):
  if len(feature_map.size()) == 3:
    feature_map = feature_map.unsqueeze(0)
  device = feature_map.device
  batch_size = feature_map.size(0)
  feature_dim = feature_map.size(1)
  size = feature_map.size(2)

  norm_batch_distance_map = torch.zeros(batch_size,size,size,device=device)
  for i in range(keypoints_size):
      vector = mean_vector_list[i]
      reshaped_feature_map = feature_map.permute(0, 2, 3, 1).view(feature_map.size(0), -1, feature_dim)
      
      batch_distance_map = F.pairwise_distance(reshaped_feature_map, vector).view(feature_map.size(0), size, size)
      batch_distance_map = 1/torch.cosh( 40*(batch_distance_map-batch_distance_map.min())
                                              /(batch_distance_map.max()-batch_distance_map.min()) )**2
      # 正規化
      m = batch_distance_map/batch_distance_map.max(1).values.max(1).values.unsqueeze(0).unsqueeze(0).repeat(112,112,1).permute(2,0,1)
      norm_batch_distance_map += m
  # 1以上を消す
  norm_batch_distance_map = (-F.relu(-norm_batch_distance_map+1)+1)
  keypoint_maps = norm_batch_distance_map.detach().cpu()

  alpha = 0.8  # Transparency factor for the heatmap overlay
  blended_tensors = (1 - alpha) * imgs + alpha * torch.cat(
    (norm_batch_distance_map.unsqueeze(1), torch.zeros(batch_size,2,size,size,device=device)),
    dim=1
    )
  blended_tensors = norm_img(blended_tensors).detach().cpu()
  
  return keypoint_maps, blended_tensors