from sklearn.manifold import TSNE import plotly.graph_objects as go import torch import torch.nn as nn import numpy as np from numba.core.errors import NumbaDeprecationWarning import warnings warnings.simplefilter('ignore', category=NumbaDeprecationWarning) from umap import UMAP from log_config import get_logger logger = get_logger(__name__) def visualize_tsne_umap_mos(model: nn.Module, dataloader: torch.utils.data.DataLoader, tsne_args: dict, umap_args: dict, device: torch.device) -> dict[go.Figure]: """ Visualize the features extracted by the model using t-SNE and UMAP with the corresponding distortion levels and MOS scores. Supports only the KADID10K dataset. Args: model (torch.nn.Module): the model to extract embeddings from dataloader (torch.utils.data.DataLoader): the data loader to use tsne_args (dict): the arguments for t-SNE umap_args (dict): the arguments for UMAP device (torch.device): the device to use for training Returns: dict with keys: - T-SNE (go.Figure): Distortions visualization with t-SNE features - UMAP (go.Figure): Distortions visualization with UMAP features - MOS_T-SNE (go.Figure): MOS visualization with t-SNE features - MOS_UMAP (go.Figure): MOS visualization with UMAP features """ logger.info('Generating visualizations...') methods = ["T-SNE", "UMAP"] # Define the colors and shades for each distortion type and level dist_color_maps = {"blur": "0, 80, 239", "color_distortion": "227, 200, 0", "jpeg": "216, 9, 168", "noise": "245, 0, 56", "brightness_change": "0, 204, 204", "spatial_distortion": "160, 82, 45", "sharpness_contrast": "96, 169, 23"} dist_shades = {1: 0.2, 2: 0.4, 3: 0.6, 4: 0.8, 5: 1.0} mos_color_map = "138, 10, 10" mos_shades = {1: 0.2, 2: 0.4, 3: 0.6, 4: 0.8, 5: 1.0} # Extract features and labels from the validation dataloader features = torch.zeros((0, model.encoder.feat_dim)) dist_colors = [] mos_colors = [] for i, batch in enumerate(dataloader, 0): img = batch["img"].to(device=device, non_blocking=True) img = img[:, 0] # Consider only the center crop dist_group = batch["dist_group"] dist_level = batch["dist_level"] mos_score = batch["mos"] with torch.cuda.amp.autocast(), torch.no_grad(): output_features, _ = model(img) features = torch.cat((features, output_features.float().detach().cpu()), 0) # Determine the color and shade for the marker based on the class and level for d_group, d_level, mos in zip(dist_group, dist_level, mos_score): mos_colors.append(f"rgba({mos_color_map}, {mos_shades[round(mos.item())]})") marker_color = dist_color_maps[d_group] marker_shade = dist_shades[d_level.item()] dist_colors.append(f"rgba({marker_color}, {marker_shade})") dist_colors = np.array(dist_colors) mos_colors = np.array(mos_colors) features = features.numpy() # Generate the figures figures = {} for method in methods: args = tsne_args if method == "T-SNE" else umap_args if method == "UMAP": features_embedded = UMAP(**args).fit_transform(features) elif method == "T-SNE": features_embedded = TSNE(**args).fit_transform(features) else: raise NotImplementedError(f"Method {method} not implemented") ### Degradations visualization fig = go.Figure() # Add legend for d_group, color in dist_color_maps.items(): placeholder = features_embedded[np.where(dist_colors == f"rgba({color}, 1.0)")[0][0]] # Needed to create the legend correctly. Store the coordinates of a point with max opacity for each class if args["n_components"] == 2: trace = go.Scatter(x=[placeholder[0]], y=[placeholder[1]], mode='markers', marker=dict(size=5, color=f"rgba({color}, 1.0)"), name=d_group) elif args["n_components"] == 3: trace = go.Scatter3d(x=[placeholder[0]], y=[placeholder[1]], z=[placeholder[2]], mode='markers', marker=dict(size=5, color=f"rgba({color}, 1.0)"), name=d_group) else: raise ValueError(f"n_components parameter must be in [2, 3].") fig.add_trace(trace) if args["n_components"] == 2: # Create 2D scatter plot with features fig.add_trace(go.Scatter(x=features_embedded[:, 0], y=features_embedded[:, 1], mode='markers', marker=dict(size=5, color=dist_colors), showlegend=False)) elif args["n_components"] == 3: # Create 3D scatter plot with features fig.add_trace(go.Scatter3d(x=features_embedded[:, 0], y=features_embedded[:, 1], z=features_embedded[:, 2], mode='markers', marker=dict(size=5, color=dist_colors), showlegend=False)) # Update layout fig.update_layout(title=f'{method} KADID10K visualization', scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), legend=dict(x=0, y=1, bgcolor='rgba(0,0,0,0)')) figures[method] = fig ### MOS visualization fig = go.Figure() # Add legend for mos, opacity in mos_shades.items(): placeholder = features_embedded[np.where(mos_colors == f"rgba({mos_color_map}, {opacity})")[0][0]] # Needed to create the legend correctly. Store the coordinates of a point for each class if args["n_components"] == 2: trace = go.Scatter(x=[placeholder[0]], y=[placeholder[1]], mode='markers', marker=dict(size=5, color=f"rgba({mos_color_map}, {opacity})"), name=mos) elif args["n_components"] == 3: trace = go.Scatter3d(x=[placeholder[0]], y=[placeholder[1]], z=[placeholder[2]], mode='markers', marker=dict(size=5, color=f"rgba({mos_color_map}, {opacity})"), name=mos) else: raise ValueError(f"n_components parameter must be in [2, 3].") fig.add_trace(trace) if args["n_components"] == 2: # Create 2D scatter plot with features fig.add_trace(go.Scatter(x=features_embedded[:, 0], y=features_embedded[:, 1], mode='markers', marker=dict(size=5, color=mos_colors), showlegend=False)) elif args["n_components"] == 3: # Create 3D scatter plot with features fig.add_trace(go.Scatter3d(x=features_embedded[:, 0], y=features_embedded[:, 1], z=features_embedded[:, 2], mode='markers', marker=dict(size=5, color=mos_colors), showlegend=False)) # Update layout fig.update_layout(title=f'{method} MOS KADID10K visualization', scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), legend=dict(x=0, y=1, bgcolor='rgba(0,0,0,0)')) figures[f"MOS_{method}"] = fig return figures