Spaces:
Sleeping
Sleeping
File size: 7,689 Bytes
7ee2717 0705c62 7ee2717 0705c62 7ee2717 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | 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
|