diff --git a/convert-annotations.py b/convert-annotations.py
index f907281..488f580 100644
--- a/convert-annotations.py
+++ b/convert-annotations.py
@@ -105,7 +105,7 @@ def read_bbox_oriented(width, height, s):
         if c_x > 1 or c_y > 1 or c_w > 1 or c_h > 1:
             print('Values are out of bounds')
         else:
-            yolo_lines.append(f'1 {c_x} {c_y} {c_w} {c_h}\n')
+            yolo_lines.append(f'2 {c_x} {c_y} {c_w} {c_h}\n')
     return yolo_lines
 
 
diff --git a/dataset-visualiser.py b/dataset-visualiser.py
new file mode 100644
index 0000000..8ec1a69
--- /dev/null
+++ b/dataset-visualiser.py
@@ -0,0 +1,28 @@
+import os
+from pathlib import Path
+import cv2
+from dto.annotationClass import AnnotationClass
+from dto.imageLabel import ImageLabel
+from preprocessing import read_labels
+from matplotlib import pyplot as plt
+
+
+annotation_classes = AnnotationClass.read_json()
+dataset_dir = os.path.join('datasets', 'zombobase-2024-06-18', 'train')
+images_dir = os.path.join(dataset_dir, 'images')
+labels_dir = os.path.join(dataset_dir, 'labels')
+
+for f in os.listdir(images_dir)[35247:]:
+    image_path = os.path.join(images_dir, f)
+    labels_path = os.path.join(labels_dir, f'{Path(f).stem}.txt')
+    img = ImageLabel(
+        image_path=image_path,
+        image=cv2.imread(image_path),
+        labels_path=labels_path,
+        labels=read_labels(labels_path)
+    )
+    img.visualize(annotation_classes)
+    print(f'visualizing {image_path}')
+    plt.close()
+    key = input('Press any key to continue')
+
diff --git a/dto/imageLabel.py b/dto/imageLabel.py
index 798052f..6aa3f8d 100644
--- a/dto/imageLabel.py
+++ b/dto/imageLabel.py
@@ -29,4 +29,4 @@ class ImageLabel:
         plt.figure(figsize=(12, 12))
         plt.axis('off')
         plt.imshow(img)
-        plt.show()
\ No newline at end of file
+        plt.show()
diff --git a/preprocessing.py b/preprocessing.py
index 4fd3010..8bc6cc0 100644
--- a/preprocessing.py
+++ b/preprocessing.py
@@ -163,11 +163,17 @@ def fix_class(folder):
         if label.startswith('0000'):
             with open(os.path.join(folder, label), 'r+') as f:
                 lines = f.readlines()
+                truncated = False
                 for i in range(0, len(lines)):
-                    l = lines[i]
-                    lines[i] = f'2{l[1:]}'
-                f.seek(0)  # rewind
-                f.writelines(lines)
+                    if len(lines[i]) < 25:
+                        print(lines[i])
+                        truncated = True
+                        lines.pop(i)
+                if truncated:
+                    f.truncate(0)
+                    f.seek(0)
+                    f.writelines(lines)
+                f.close()
 
 
 if __name__ == '__main__':
diff --git a/train.py b/train.py
index 2cd2daa..566ec3c 100644
--- a/train.py
+++ b/train.py
@@ -102,6 +102,7 @@ if __name__ == '__main__':
 
     model = YOLOv10(latest_model or 'yolov10x.yaml')
 
+    folder = f'{prefix}2024-06-18'
     yaml = abspath(path.join('datasets', today_folder, 'data.yaml'))
     results = model.train(data=yaml, epochs=100, batch=10, imgsz=640)