testing new latent space visualization code

ContraCLIP/traverse_latent_space.py

@@ -10,6 +10,10 @@ from torchvision.transforms import ToPILImage
 from lib import SupportSets, GENFORCE_MODELS, update_progress, update_stdout, STYLEGAN_LAYERS
 from models.load_generator import load_generator
 
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from sklearn.manifold import TSNE
 
 class DataParallelPassthrough(nn.DataParallel):
     def __getattr__(self, name):
@@ -97,6 +101,63 @@ def create_gif(image_list, gif_height=256):
 
     return transformed_images_gif_frames
 
+def visualize_latent_space(tsne_latent_codes, semantic_dipoles, output_dir, save_filename="latent_space_tsne.png", shift_steps=16):
+    """
+    Visualize the t-SNE reduced latent space with minimal annotations.
+
+    Args:
+        tsne_latent_codes (np.ndarray): The 3D latent codes after t-SNE transformation.
+        semantic_dipoles (list): List of semantic directions (labels) for paths.
+        shift_steps (int): Number of positive/negative steps along each path.
+        output_dir (str): Directory to save the generated plot.
+        save_filename (str): Name of the file to save the plot.
+    """
+    fig = plt.figure(figsize=(16, 12))  # Larger figure for clarity
+    ax = fig.add_subplot(111, projection='3d')
+
+    num_paths = len(semantic_dipoles)  # Each dipole represents one path
+    cmap = plt.cm.get_cmap('tab10', num_paths)
+
+    for i in range(num_paths):
+        # Indices for the path in tsne_latent_codes
+        start_idx = i * (2 * shift_steps + 1)
+        pos_idx = start_idx + shift_steps  # Positive endpoint
+        neg_idx = start_idx + 2 * shift_steps  # Negative endpoint
+
+        # Extract path points
+        path_indices = list(range(start_idx, neg_idx + 1))
+        path_coords = tsne_latent_codes[path_indices]
+
+        # Plot the entire path (all intermediate points in a single color)
+        ax.plot(
+            path_coords[:, 0], path_coords[:, 1], path_coords[:, 2],
+            color=cmap(i),
+            linewidth=2
+        )
+
+        # Extract positive and negative endpoint coordinates
+        pos_coords = tsne_latent_codes[pos_idx]
+        neg_coords = tsne_latent_codes[neg_idx]
+
+        # Plot positive and negative endpoints
+        ax.scatter(*pos_coords, color=cmap(i), s=100, label=f"{semantic_dipoles[i][0]} → {semantic_dipoles[i][1]}")
+        ax.scatter(*neg_coords, color=cmap(i), s=100)
+
+    # Add legend
+    ax.legend(loc='best', fontsize=10)
+
+    # Set titles and labels
+    ax.set_title("t-SNE Latent Space Visualization")
+    ax.set_xlabel("t-SNE Dimension 1")
+    ax.set_ylabel("t-SNE Dimension 2")
+    ax.set_zlabel("t-SNE Dimension 3")
+
+    # Save the plot
+    os.makedirs(output_dir, exist_ok=True)
+    save_path = osp.join(output_dir, save_filename)
+    plt.savefig(save_path, dpi=300, bbox_inches="tight")
+    print(f"Visualization saved to {save_path}")
+
 
 def main():
     """ContraCLIP -- Latent space traversal script.
@@ -210,6 +271,8 @@ def main():
     # -- Get prompt corpus list
     with open(osp.join(models_dir, 'semantic_dipoles.json'), 'r') as f:
         semantic_dipoles = json.load(f)
+    
+    # semantic_directions = [f"{dipole[0]} → {dipole[1]}" for dipole in semantic_dipoles]
 
     # Check given pool directory
     pool = osp.join('experiments', 'latent_codes', gan, args.pool)
@@ -321,6 +384,9 @@ def main():
         print("  \\__Shift steps                 : {}".format(2 * args.shift_steps))
         print("  \\__Traversal length            : {}".format(round(2 * args.shift_steps * args.eps, 3)))
 
+    # Store latent codes for T-SNE visualization (for all paths across each latent code)
+    all_paths_latent_codes = []
+
     # Iterate over given latent codes
     for i in range(num_of_latent_codes):
         # Get latent code
@@ -333,6 +399,9 @@ def main():
                                                                                   num_of_latent_codes),
                             num_of_latent_codes, i)
 
+        # Append the starting latent code to tsne_latent_codes
+        # tsne_latent_codes.append(x_.clone().cpu().numpy().flatten())
+
         # Create directory for current latent code
         latent_code_dir = osp.join(out_dir, '{}'.format(latent_code_hash))
         os.makedirs(latent_code_dir, exist_ok=True)
@@ -386,7 +455,7 @@ def main():
                     latent_code = latent_code[:, 0, :]
 
             cnt = 0
-            for _ in range(args.shift_steps):
+            for k in range(args.shift_steps):
                 cnt += 1
 
                 # Calculate shift vector based on current z
@@ -410,6 +479,10 @@ def main():
                     latent_code = latent_code + shift
                     current_path_latent_code = latent_code
 
+                # Append intermediate latent code
+                # if k != args.shift_steps - 1:
+                #     tsne_latent_codes.append(latent_code.clone().cpu().numpy().flatten())
+
                 # Store latent codes and shifts
                 if cnt == args.shift_leap:
                     if ('stylegan' in gan) and (stylegan_space == 'W+'):
@@ -421,6 +494,8 @@ def main():
                     current_path_latent_codes.append(current_path_latent_code)
                     cnt = 0
             positive_endpoint = latent_code.clone().reshape(1, -1)
+
+            # tsne_latent_codes.append(positive_endpoint.clone().cpu().numpy().flatten())
             # ========================
 
             # == Negative direction ==
@@ -430,7 +505,7 @@ def main():
                 if stylegan_space == 'W':
                     latent_code = latent_code[:, 0, :]
             cnt = 0
-            for _ in range(args.shift_steps):
+            for k in range(args.shift_steps):
                 cnt += 1
                 # Calculate shift vector based on current z
                 support_sets_mask = torch.zeros(1, LSS.num_support_sets)
@@ -453,6 +528,10 @@ def main():
                     latent_code = latent_code + shift
                     current_path_latent_code = latent_code
 
+                # Append intermediate latent code
+                # if k != args.shift_steps - 1:
+                #     tsne_latent_codes.append(latent_code.clone().cpu().numpy().flatten())
+
                 # Store latent codes and shifts
                 if cnt == args.shift_leap:
                     if ('stylegan' in gan) and (stylegan_space == 'W+'):
@@ -464,6 +543,8 @@ def main():
                     current_path_latent_codes = [current_path_latent_code] + current_path_latent_codes
                     cnt = 0
             negative_endpoint = latent_code.clone().reshape(1, -1)
+
+            # tsne_latent_codes.append(latent_code.clone().cpu().numpy().flatten())
             # ========================
 
             # Calculate latent path phi coefficient (end-to-end distance / latent path length)
@@ -531,13 +612,69 @@ def main():
 
         # Save all latent paths and shifts for the current latent code (sample) in a tensor of size:
         #   paths_latent_codes : torch.Size([num_gen_paths, 2 * args.shift_steps + 1, G.dim_z])
-        torch.save(torch.cat(paths_latent_codes), osp.join(latent_code_dir, 'paths_latent_codes.pt'))
+        paths_latent_codes_tensor = torch.cat(paths_latent_codes)
+        torch.save(paths_latent_codes_tensor, osp.join(latent_code_dir, 'paths_latent_codes.pt'))
+        all_paths_latent_codes.append(paths_latent_codes_tensor.cpu().numpy())
 
         if args.verbose:
             update_stdout(1)
             print()
             print()
 
+    # After processing all latent codes and paths
+    if args.verbose:
+        print("Performing t-SNE on latent codes for visualization...")
+
+    # # Consolidate all paths for T-SNE visualization (total_paths = num_of_latent_codes * num_gen_paths)
+    # all_paths_np = np.concatenate(all_paths_latent_codes, axis=0)  # Shape: [total_paths, steps_per_path, latent_dim]
+    # all_paths_flattened = all_paths_np.reshape(-1, all_paths_np.shape[-1])  # Flatten paths into 2D array for T-SNE
+
+    # # Apply 3D T-SNE
+    # tsne_model = TSNE(n_components=3, perplexity=30, learning_rate=200, random_state=42)
+    # tsne_transformed = tsne_model.fit_transform(all_paths_flattened)  # Shape: [total_points, 3]
+
+    # path_indices = []  # List to store indices for each path
+    # start_idx = 0  # Starting index for the current path in all_paths_np
+
+    # steps_per_path = 2 * args.shift_steps + 1  # Number of points in each path
+
+    # # Iterate over each latent code and its paths
+    # for i in range(num_of_latent_codes):  # Loop through latent codes
+    #     for dim in range(num_gen_paths):  # Loop through directions (paths)
+    #         # Generate the indices for this path
+    #         indices = list(range(start_idx, start_idx + steps_per_path))
+    #         path_indices.append(indices)
+            
+    #         # Update the starting index for the next path
+    #         start_idx += steps_per_path
+
+
+    all_paths_latent_code_0 = all_paths_latent_codes[0]
+    num_paths, num_steps, _ = all_paths_latent_code_0.shape
+    tsne_latent_codes = all_paths_latent_code_0.reshape(-1, all_paths_latent_code_0.shape[-1])
+
+    # Apply 3D T-SNE
+    tsne_model = TSNE(n_components=3, perplexity=30, learning_rate=200, random_state=42)
+    tsne_transformed = tsne_model.fit_transform(tsne_latent_codes)  # Shape: [total_points = num_paths * num_steps, 3]
+
+    # For this specific latent code, generate indices for each of its paths
+    path_indices = []
+    start_idx = 0
+    for _ in range(num_paths):
+        indices = list(range(start_idx, start_idx + num_steps))
+        path_indices.append(indices)
+        start_idx += num_steps
+
+    
+    tsne_vis_dir = osp.join(out_dir, 'tsne_visualizations')
+    visualize_latent_space(
+        tsne_latent_codes=tsne_transformed,  # T-SNE-reduced latent codes
+        semantic_dipoles=semantic_dipoles,  # Semantic labels for paths
+        paths=path_indices,  # Indices of paths (for a single latent code)
+        output_dir=tsne_vis_dir,
+        save_filename="latent_space_tsne.png"
+    )
+    
     # Create summarizing MD files
     if args.gif or args.strip:
         # For each interpretable path (warping function), collect the generated image sequences for each original latent