Update visualization.py

visualization.py

@@ -153,9 +153,13 @@ def plot_mse_heatmap(mse_values, title, df):
     sns.heatmap(mse_2d, cmap='YlOrRd', cbar=False, ax=ax)
 
     # Set x-axis ticks to timecodes
-    num_ticks = 60
+    num_ticks = min(60, len(mse_values))
     tick_locations = np.linspace(0, len(mse_values) - 1, num_ticks).astype(int)
-    tick_labels = [df['Timecode'].iloc[i] for i in tick_locations]
+    
+    # Ensure tick_locations are within bounds
+    tick_locations = tick_locations[tick_locations < len(df)]
+    
+    tick_labels = [df['Timecode'].iloc[i] if i < len(df) else '' for i in tick_locations]
 
     ax.set_xticks(tick_locations)
     ax.set_xticklabels(tick_labels, rotation=90, ha='center', va='top')