diff --git a/curvature_analysis.py b/curvature_analysis.py
index eaf5bea..7c875d9 100644
--- a/curvature_analysis.py
+++ b/curvature_analysis.py
@@ -132,7 +132,7 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
         frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
     frame_labels[curvature >= bins[-2]] = labels[-1]
        
-    return frame_labels
+    return frame_labels, curvature
 
 import tensorflow as tf
 from configs.config_abl_overfit_jumpsuit_2 import config
@@ -185,14 +185,15 @@ def visualize_curvature_labels(folder, results_folder, frame):
 
     uv_curvature = np.load(os.path.join(folder, 'uv_curvature', f'{id}_{frame_nr}.npz'))['curvature_uv']
     mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', f'{id}_{frame_nr}.npz'))['mesh_curvature'][mask]
+    mesh_curvature_raw = np.load(os.path.join(folder, 'mesh_curvature_raw', f'{id}_{frame_nr}.npz'))['mesh_curvature_raw'][mask]
     uv_curvature = tf.one_hot(uv_curvature, 7)
     pred_new = get_prediction(folder, frame, uv_curvature)
     pred_new = postprocess(pred_new, faces)
 
-    curvature_labels_gt = generate_curvatures(verts, faces, mask, 25)
-    curvature_labels_pred = generate_curvatures(pred, faces, mask, 25)
+    curvature_labels_gt, curv_gt = generate_curvatures(verts, faces, mask, 25)
+    curvature_labels_pred, _ = generate_curvatures(pred, faces, mask, 25)
     curvatures_new = (uv2meshlabel(uv_curvature[None,:,:,:])[0].numpy() * mask[:,None])[mask, :]
-    curvatures_pred_new = generate_curvatures(pred_new, faces, mask, 25)
+    curvatures_pred_new, curv_gt_new = generate_curvatures(pred_new, faces, mask, 25)
     dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:])
     dice_new = dice_dense(tf.one_hot(curvatures_pred_new, 7)[None,:], curvatures_new[None,:])
     curvatures_new = curvatures_new.argmax(axis=-1) - 1