Update generate_label.py

lab-midas · Jul 12, 2021 · f5695c0 · f5695c0
1 parent 40bbb48
commit f5695c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/med_io/generate_label.py b/med_io/generate_label.py
@@ -30,7 +30,7 @@ def make_class_label(label, patch_pos, patch_size):
             new_label = label*-1+(patch_pos[idx][2] + patch_size[2] // 2)
             class_pos = tf.where(new_label > 0)[0][-1]
 
-            one_hot = tf.one_hot(class_pos, config['body_identification_n_classes'], dtype=tf.float32)
+            one_hot = tf.one_hot(class_pos, len(config['name_output_channel'][config['dataset']]), dtype=tf.float32)
             temp[idx]=one_hot