diff --git a/curvature_analysis.py b/curvature_analysis.py
index 7bfe55a..0b7e952 100644
--- a/curvature_analysis.py
+++ b/curvature_analysis.py
@@ -133,6 +133,7 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
        
     return frame_labels
 
+import tensorflow as tf
 
 def visualize_curvature_labels(folder, results_folder, frame):
     id = frame.split('_')[0]
@@ -146,7 +147,8 @@ def visualize_curvature_labels(folder, results_folder, frame):
 
     curvature_labels_gt = generate_curvatures(verts, faces, mask, 25)
     curvature_labels_pred = generate_curvatures(pred, faces, mask, 25)
-    dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
+    # dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
+    dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:]).numpy()
     colors_binned_all = []
 
     for curvature_labels in [curvature_labels_gt, curvature_labels_pred]: