add dataset-visualiser.py

This commit is contained in:
Alex Bezdieniezhnykh
2024-06-19 02:28:50 +03:00
parent 165d609d01
commit 6c50dd19b7
5 changed files with 41 additions and 6 deletions
+1 -1
View File
@@ -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
+28
View File
@@ -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')
+10 -4
View File
@@ -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__':
+1
View File
@@ -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)