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

class GetTransformedCoords(nn.Module):
  def __init__(self, affine_matrix, center):
    super().__init__()
    self.affine_matrix = affine_matrix
    self.center = center

  def forward(self, _coords):
    # coords: like tensor([[43, 26], [44, 27], [45, 28]])
    center_x, center_y = self.center
    # 元の座標を中心原点にシフト
    coords = _coords.clone()
    coords[:, 0] -= center_x
    coords[:, 1] -= center_y

    # 各バッチに対して変換を行う
    homogeneous_coordinates = torch.cat([coords, torch.ones(coords.shape[0], 1, dtype=torch.float32, device=coords.device)], dim=1)
    transformed_coordinates = torch.bmm(self.affine_matrix, homogeneous_coordinates.unsqueeze(-1)).squeeze(-1)

    # 画像の範囲内に収める
    # transformed_x = max(0, min(width - 1, transformed_coordinates[:, 0]))
    # transformed_y = max(0, min(height - 1, transformed_coordinates[:, 1]))
    transformed_x = transformed_coordinates[:, 0]
    transformed_y = transformed_coordinates[:, 1]

    transformed_x += center_x
    transformed_y += center_y
    return torch.stack([transformed_x, transformed_y]).t().to(torch.long)

# ルートを取らないpairwise_distanceのバージョン
def pairwise_distance_squared(a, b):
    return torch.sum((a - b) ** 2, dim=-1)

def cosine_similarity(a, b):
  # ベクトルaとbの内積を計算
  dot_product = torch.matmul(a, b)
  # ベクトルaとbのノルム（大きさ）を計算
  norm_a = torch.sqrt(torch.sum(a ** 2, dim=-1))
  norm_b = torch.sqrt(torch.sum(b ** 2, dim=-1))
  # コサイン類似度を計算（内積をノルムの積で割る）
  return dot_product / (norm_a * norm_b)

def batch_cosine_similarity(a, b):
  # ベクトルaとbの内積を計算
  dot_product = torch.einsum('bnd,bnd->bn', a, b)
  # ベクトルaとbのノルム（大きさ）を計算
  norm_a = torch.sqrt(torch.sum(a ** 2, dim=-1))
  norm_b = torch.sqrt(torch.sum(b ** 2, dim=-1))
  # コサイン類似度を計算（内積をノルムの積で割る）
  return dot_product / (norm_a * norm_b)

class TripletLossBatch(nn.Module):
  def __init__(self):
    super(TripletLossBatch, self).__init__()

  def forward(self, anchor, positive, negative, margin=1.0):
    distance_positive = F.pairwise_distance(anchor, positive, p=2)
    distance_negative = F.pairwise_distance(anchor, negative, p=2)
    losses = torch.relu(distance_positive - distance_negative + margin)
    return losses.mean()

class TripletLossCosineSimilarity(nn.Module):
  def __init__(self):
    super(TripletLossCosineSimilarity, self).__init__()

  def forward(self, anchor, positive, negative, margin=1.0):
    distance_positive = 1 - batch_cosine_similarity(anchor, positive)
    distance_negative = 1 - batch_cosine_similarity(anchor, negative)
    losses = torch.relu(distance_positive - distance_negative + margin)
    return losses.mean()

def imsave(img):
  img = torchvision.utils.make_grid(img)
  img = img / 2 + 0.5
  npimg = img.detach().cpu().numpy()
  # plt.imshow(np.transpose(npimg, (1, 2, 0)))
  # plt.show()
  # save image
  npimg = np.transpose(npimg, (1, 2, 0))
  npimg = npimg * 255
  npimg = npimg.astype(np.uint8)
  Image.fromarray(npimg).save('sample.png')
  
def norm_img(img):
  return (img-img.min())/(img.max()-img.min())

def norm_img2(img):
  return (img-img.min())/(img.max()-img.min())*255

class DistanceMapLogger:
  def __call__(self, img, feature_map, save_path, x_coords=None, y_coords=None):
    device = feature_map.device
    batch_size = feature_map.size(0)
    feature_dim = feature_map.size(1)
    image_size = feature_map.size(2)
    
    if x_coords is None:
      x_coords = [69]*batch_size
    if y_coords is None:
      y_coords = [42]*batch_size

    # PCAで3次元のマップを抽出
    pca = PCA(n_components=3)
    pca_result = pca.fit_transform(feature_map.permute(0,2,3,1).reshape(-1,feature_dim).detach().cpu().numpy())  # PCA を実行
    reshaped_pca_result = pca_result.reshape(batch_size,image_size,image_size,3) # 3次元に変換（元は1次元）

    sample_num = 0
    vectors = feature_map[torch.arange(feature_map.size(0)), :, y_coords, x_coords]  # 1次元ベクトルに合わせてサイズを調整
    vector = vectors[sample_num]

    # バッチ内の各特徴マップに対して内積を計算
    # feature_mapの次元を並べ替えてバッチと高さ・幅を平坦化
    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)
    # batch_distance_map = F.cosine_similarity(reshaped_feature_map, vector.unsqueeze(0).unsqueeze(0).expand(65,size*size,32), dim=2).permute(1, 0).reshape(feature_map.size(0), size, 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
    # norm_batch_distance_map[:,0,0] = 0.001
    # 可視化と保存
    fig, axes = plt.subplots(5, 4, figsize=(20, 25))
    for ax in axes.flatten():
        ax.axis('off')
    # 余白をなくす
    plt.subplots_adjust(wspace=0, hspace=0)
    # 外の余白もなくす
    plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
    
    # 距離マップの可視化
    for i in range(5):
      axes[i, 0].imshow(norm_batch_distance_map[i].detach().cpu(), cmap='hot')
      if i == sample_num:
        axes[i, 0].scatter(x_coords[i], y_coords[i], c='b', s=7)

      distance_map = torch.cat(((norm_batch_distance_map[i]/norm_batch_distance_map[i].max()).unsqueeze(0),torch.zeros(2,image_size,image_size,device=device)))
      alpha = 0.9  # Transparency factor for the heatmap overlay
      blended_tensor = (1 - alpha) * img[i] + alpha * distance_map
      axes[i, 1].imshow(norm_img(blended_tensor.permute(1,2,0).detach().cpu()))

      axes[i, 2].imshow(norm_img(img[i].permute(1,2,0).detach().cpu()))
    
      axes[i, 3].imshow(norm_img(reshaped_pca_result[i]))
      
    plt.savefig(save_path)



def get_heatmaps(self, img, feature_map, source_num=0, target_num=1,  x_coords=69, y_coords=42):
  device = feature_map.device
  batch_size = feature_map.size(0)
  feature_dim = feature_map.size(1)
  image_size = feature_map.size(2)

  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]  # 1次元ベクトルに合わせてサイズを調整
  vector = vectors[source_num]

  # バッチ内の各特徴マップに対して内積を計算
  # feature_mapの次元を並べ替えてバッチと高さ・幅を平坦化
  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)
  # batch_distance_map = F.cosine_similarity(reshaped_feature_map, vector.unsqueeze(0).unsqueeze(0).expand(65,size*size,32), dim=2).permute(1, 0).reshape(feature_map.size(0), size, 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
  # norm_batch_distance_map[:,0,0] = 0.001

  source_map = norm_batch_distance_map[source_num]
  target_map = norm_batch_distance_map[target_num]

  alpha = 0.9
  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)))
  
  return source_map, target_map, blended_source, blended_target