refactoring

app.py

@@ -6,7 +6,6 @@ from torch.utils.data import DataLoader
 import matplotlib.pyplot as plt
 from model_module import AutoencoderModule
 from dataset import MyDataset, load_filenames
-from utils import DistanceMapLogger
 import numpy as np
 from PIL import Image
 import base64
@@ -159,7 +158,6 @@ async () => {
             document.querySelector("#input_file").click();
         };
 
-        // GradioのFileコンポーネントから画像を読み込む
         document.querySelector("#input_file").addEventListener("change", function(e) {
             const files = e.target.files;
             console.log(files);
@@ -182,7 +180,6 @@ async () => {
             }
         });
 
-        // GradioボタンにJavaScriptの機能を追加
         document.getElementById("crop_button").onclick = function() {
             if (cropper) {
                 const canvas = cropper.getCroppedCanvas();
@@ -227,13 +224,9 @@ with gr.Blocks() as demo:
             gr.HTML('<input type="file" id="input_file" style="display:none;">')
             input_file_button = gr.Button("Upload and Crop Image", elem_id="input_file_button", variant="primary")
             crop_button = gr.Button("Crop", elem_id="crop_button", variant="primary")
-            # 画像を表示するためのHTML画像タグをGradioで表示
             gr.HTML('<img id="crop_view" style="max-width:100%;">')
-            # Gradioのボタンコンポーネントを追加し、IDを付与
-            # クロップされた画像データのテキストボックス（Base64データ）
             cropped_image_data = gr.Textbox(visible=False, elem_id="cropped_image_data")
-            input_image = gr.Image(label="Cropped Image", elem_id="input_image")
-            # cropped_image_dataが更新されたらprocess_imageを呼び出す
+            input_image = gr.Image(visible=False, label="Cropped Image", elem_id="input_image")
             cropped_image_data.change(process_image, inputs=cropped_image_data, outputs=input_image)
 
             # examples
@@ -248,15 +241,17 @@ with gr.Blocks() as demo:
         with gr.Column():
             output_plot = gr.Plot()
 
+    # ヒートマップの更新
+    source_num.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
+    x_coords.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
+    y_coords.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
 
-        # Gradioインターフェースの代わり
-        source_num.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
-        x_coords.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
-        y_coords.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
-        input_image.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
+    def change_input_image_hanlder(source_num, x_coords, y_coords, input_image):
+        visible = False if input_image is None else True
+        return get_heatmaps(source_num, x_coords, y_coords, input_image), gr.Image(visible=visible, label="Cropped Image", elem_id="input_image")
+    input_image.change(change_input_image_hanlder, inputs=[source_num, x_coords, y_coords, input_image], outputs=[output_plot, input_image])
 
-        # JavaScriptコードをロード
-        demo.load(None, None, None, js=scripts)        
-        
+    # JavaScriptコードをロード
+    demo.load(None, None, None, js=scripts)
     demo.launch()
 

dataset.py

@@ -8,27 +8,6 @@ import os
 from utils import RandomAffineAndRetMat
 
 def load_filenames(data_dir):
-  # label_data = pd.read_json(INPUT_DIR+'DataList.json')
-  # label_data = label_data.sort_index()
-  # tmp_points = []
-  # filenames = []
-
-  # for o in tqdm(label_data.data[0:1000]):
-  #   filenames.append(o['filename'])
-  #   a = o['filename']
-
-  #   tmps = []
-  #   for i in range(60):
-  #     tmps.append(o['points'][str(i)]['x'])
-  #     tmps.append(o['points'][str(i)]['y'])
-  #   tmp_points.append(tmps) # datanum
-
-  # filenames = pd.Series(filenames)
-  # filenames = [str(i).zfill(4)+'.jpg' for i in range(3400)]
-  # df_points = pd.DataFrame(tmp_points)
-  
-  # load from data_dir
-  # 画像の拡張子のみ
   img_exts = ['.jpg', '.jpeg', '.png', '.bmp', '.ppm', '.pgm', '.tif', '.tiff']
   filenames = [f for f in os.listdir(data_dir) if os.path.splitext(f)[1].lower() in img_exts]
 

model_module.py

@@ -1,145 +1,11 @@
 import pytorch_lightning as pl
-import torch
-from torch import nn
-from torch.optim import SGD
-from torchvision.utils import save_image
-import os
-from utils import TripletLossBatch, pairwise_distance_squared, GetTransformedCoords, DistanceMapLogger
 from model import Autoencoder
 
 class AutoencoderModule(pl.LightningModule):
-    def __init__(self, feature_dim=64, learning_rate=0.1, lambda_c=0.97, initial_margin=1.0, initial_threshold=2.0, save_interval=100, output_dir="output_images"):
+    def __init__(self, feature_dim=64):
         super(AutoencoderModule, self).__init__()
         self.feature_dim = feature_dim
-        self.learning_rate = learning_rate
-        self.lambda_c = lambda_c
-        self.margin_img = initial_margin
-        self.margin_img_init = initial_margin
-        self.threshold = initial_threshold
         self.model = Autoencoder(self.feature_dim)
-        self.criterion = nn.MSELoss()
-        self.triplet_loss = TripletLossBatch()
-        self.losses = []
-        self.save_interval = save_interval  # バッチごとの出力間隔
-        self.output_dir = output_dir
-        os.makedirs(self.output_dir, exist_ok=True)
 
     def forward(self, x):
         return self.model(x)
-
-    def training_step(self, batch, batch_idx):
-        img, mat, _ = batch
-        batch_size, _, _, size, size = img.shape
-        img = img.view(batch_size*2, 3, size, size)
-        mat = mat.view(batch_size*2, 3, 3)
-        
-        dec5_output, output = self.model(img)
-        mse_loss = self.criterion(output, img)
-        
-        # 画像内方向の処理
-        num_anchor_sets = 2**12
-        trip_loss = 0
-        std_list = [2.5*1.025**self.current_epoch, 5*1.025**self.current_epoch]
-        for c in std_list:
-            std = size / c
-            anchors = torch.randint(0, size, (batch_size*2, num_anchor_sets, 1, 2))
-            coords = anchors + torch.normal(0, std, (batch_size*2, num_anchor_sets, 2, 2)).long()
-            valid_coords_idx = (((coords >= 0) & (coords < size)).sum(3) == 2).sum(2) != 2
-            coords[valid_coords_idx] = 0
-            anchors[valid_coords_idx] = 0
-            
-            # 最も近い座標の選択
-            d = pairwise_distance_squared(anchors.float(), coords.float())
-            idx = torch.argmin(d, dim=2)
-            anchors, positives, negatives = self._get_triplet_coordinates(anchors, coords, idx)
-            
-            # dec5_outputから特徴ベクトルを抽出
-            anchor_vectors, positive_vectors, negative_vectors = self._extract_feature_vectors(dec5_output, batch_size, anchors, positives, negatives)
-            trip_loss += self.triplet_loss(anchor_vectors, positive_vectors, negative_vectors, self.margin_img)
-        
-        trip_loss /= len(std_list)
-        self.margin_img = self.margin_img_init + self.margin_img - trip_loss.detach()
-        
-        # 変形の学習
-        num_samples = 2**20
-        tf_loss = self._compute_transformation_loss(dec5_output, mat, batch_size, size, num_samples)
-        
-        # バッチ方向の処理
-        bat_dist_loss = self._compute_batch_direction_loss(dec5_output, batch_size, size)
-        
-        # 合計損失
-        loss = mse_loss + trip_loss + 0.001 * bat_dist_loss + (0.001 * 1.**self.current_epoch) * tf_loss
-        self.log("train_loss", loss)
-
-        # VRAM管理
-        del img, output
-        torch.cuda.empty_cache()
-        
-        return loss
-
-
-    def _get_triplet_coordinates(self, anchors, coords, idx):
-        anchors = anchors.squeeze(2)
-        positives = coords[torch.arange(coords.size(0))[:, None, None], torch.arange(coords.size(1))[None, :, None], idx[:, :, None]].squeeze(2)
-        negatives = coords[torch.arange(coords.size(0))[:, None, None], torch.arange(coords.size(1))[None, :, None], (1 - idx)[:, :, None]].squeeze(2)
-        return anchors, positives, negatives
-
-    def _extract_feature_vectors(self, dec5_output, batch_size, anchors, positives, negatives):
-        y_anchors = anchors[:, :, 0].unsqueeze(2).expand(-1, -1, self.feature_dim)
-        x_anchors = anchors[:, :, 1].unsqueeze(2).expand(-1, -1, self.feature_dim)
-        y_positives = positives[:, :, 0].unsqueeze(2).expand(-1, -1, self.feature_dim)
-        x_positives = positives[:, :, 1].unsqueeze(2).expand(-1, -1, self.feature_dim)
-        y_negatives = negatives[:, :, 0].unsqueeze(2).expand(-1, -1, self.feature_dim)
-        x_negatives = negatives[:, :, 1].unsqueeze(2).expand(-1, -1, self.feature_dim)
-
-        anchor_vectors = dec5_output[torch.arange(batch_size*2)[:, None, None], torch.arange(self.feature_dim), y_anchors, x_anchors]
-        positive_vectors = dec5_output[torch.arange(batch_size*2)[:, None, None], torch.arange(self.feature_dim), y_positives, x_positives]
-        negative_vectors = dec5_output[torch.arange(batch_size*2)[:, None, None], torch.arange(self.feature_dim), y_negatives, x_negatives]
-        return anchor_vectors, positive_vectors, negative_vectors
-
-    def _compute_transformation_loss(self, dec5_output, mat, batch_size, size, num_samples=2**12):
-        anchor_indices = torch.randint(batch_size, (num_samples, 1), device=self.device).repeat(1, 2).reshape(num_samples*2)
-        coords_x = torch.randint(0, size, (num_samples, 1), dtype=torch.float32, device=self.device).repeat(1, 2).reshape(num_samples*2, 1)
-        coords_y = torch.randint(0, size, (num_samples, 1), dtype=torch.float32, device=self.device).repeat(1, 2).reshape(num_samples*2, 1)
-        anchor_coords = torch.cat((coords_x, coords_y), 1)
-        anchor_mat = mat[anchor_indices]
-        tf_anchor_coords = GetTransformedCoords(anchor_mat, [size/2, size/2])(anchor_coords)
-
-        anchor_vectors = torch.zeros([num_samples*2, self.feature_dim], device=self.device)
-        inner_idx_flat = ((0 <= tf_anchor_coords[:,0]) & (tf_anchor_coords[:,0] < size)) & ((0 <= tf_anchor_coords[:,1]) & (tf_anchor_coords[:,1] < size))
-        anchor_vectors[inner_idx_flat] = dec5_output[anchor_indices[inner_idx_flat], :, tf_anchor_coords[inner_idx_flat, 0], tf_anchor_coords[inner_idx_flat, 1]]
-
-        inner_idx_and = inner_idx_flat.view(num_samples, 2).t()[0] & inner_idx_flat.view(num_samples, 2).t()[1]
-        anchor_vectors = anchor_vectors.view(num_samples, 2, self.feature_dim)[inner_idx_and]
-        return pairwise_distance_squared(anchor_vectors[:,0], anchor_vectors[:,1]).mean()
-
-    def _compute_batch_direction_loss(self, dec5_output, batch_size, size):
-        N = 2**12
-        anchor_indices = torch.randint(0, batch_size, (N,)) * 2 + torch.randint(0, 2, (N,))
-        anchor_coords = torch.randint(0, size, (N, 2))
-        other_indices = torch.randint(0, batch_size-1, (N, 2)) * 2 + torch.randint(0, 2, (N, 2))
-        other_indices += (other_indices >= anchor_indices.unsqueeze(1)).long() * 2
-        other_coords = torch.randint(0, size, (N, 2, 2))
-
-        anchor_vectors = dec5_output[anchor_indices, :, anchor_coords[:, 0], anchor_coords[:, 1]]
-        other_vectors = dec5_output[other_indices, :, other_coords[:, :, 0], other_coords[:, :, 1]]
-        distances = pairwise_distance_squared(anchor_vectors.unsqueeze(1), other_vectors)
-        return distances[distances < self.threshold].sum() / ((distances < self.threshold).sum() + 1e-10)
-
-    def configure_optimizers(self):
-        optimizer = SGD(self.parameters(), lr=self.learning_rate)
-        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch: self.lambda_c**epoch)
-        return [optimizer], [scheduler]
-
-    # def save_intermediate_image(self, output, epoch):
-    #     save_image(output[:4], os.path.join(self.output_dir, f"epoch_{epoch}_output.png"), nrow=1)
-    #     print(f"Saved intermediate image at epoch {epoch}")
-        
-    # def distance_map(self, _input, feature_map, epoch, x_coords=None, y_coords=None):
-    #     save_path = os.path.join(self.output_dir, f"epoch_{epoch}_distance_map.png")
-    #     DistanceMapLogger()(_input, feature_map, save_path, x_coords, y_coords)
-
-    def configure_optimizers(self):
-        optimizer = SGD(self.parameters(), lr=self.learning_rate)
-        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch: self.lambda_c**epoch)
-        return [optimizer], [scheduler]

utils.py

@@ -77,176 +77,3 @@ class RandomAffineAndRetMat(torch.nn.Module):
     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