azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 08:36:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

add train.py

Browse Source 
form dataset for current date
add exception catching
This commit is contained in:
Oleksandr Bezdieniezhnykh
2024-06-05 23:35:06 +03:00
parent 3fd726f9c7
commit 07ea67746a
7 changed files with 168 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
preprocessing.py
+27 -18
View File
@@ -3,15 +3,11 @@ import time
from pathlib import Path
import albumentations as A
import cv2
from dto.annotationClass import AnnotationClass
from constants import current_images_dir, current_labels_dir, annotation_classes
from dto.imageLabel import ImageLabel
labels_dir = 'labels'
images_dir = 'images'
current_dataset_dir = os.path.join('datasets', 'zombobase-current')
current_images_dir = os.path.join(current_dataset_dir, 'images')
current_labels_dir = os.path.join(current_dataset_dir, 'labels')
annotation_classes = AnnotationClass.read_json()
def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
@@ -37,16 +33,19 @@ def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
results = []
for i, transform in enumerate(transforms):
res = transform(image=img_ann.image, bboxes=img_ann.labels)
path = Path(img_ann.image_path)
name = f'{path.stem}_{i+1}'
img = ImageLabel(
image=res['image'],
labels=res['bboxes'],
image_path=os.path.join(current_images_dir, f'{name}{path.suffix}'),
labels_path=os.path.join(current_labels_dir, f'{name}.txt')
)
results.append(img)
try:
res = transform(image=img_ann.image, bboxes=img_ann.labels)
path = Path(img_ann.image_path)
name = f'{path.stem}_{i+1}'
img = ImageLabel(
image=res['image'],
labels=res['bboxes'],
image_path=os.path.join(current_images_dir, f'{name}{path.suffix}'),
labels_path=os.path.join(current_labels_dir, f'{name}.txt')
)
results.append(img)
except Exception as e:
print(f'Error during transformtation: {e}')
return results
@@ -74,7 +73,7 @@ def read_labels(labels_path) -> [[]]:
for row in rows:
str_coordinates = row.split(' ')
class_num = str_coordinates.pop(0)
coordinates = [float(n) for n in str_coordinates]
coordinates = [float(n.replace(',', '.')) for n in str_coordinates]
coordinates.append(class_num)
arr.append(coordinates)
return arr
@@ -111,8 +110,18 @@ def main():
labels_path=labels_path,
labels=read_labels(labels_path)
))
except FileNotFoundError:
print(f'No labels file {labels_path} found')
except Exception as e:
print(f'Error appeared {e}')
try:
os.remove(image_path)
except OSError:
pass
try:
os.remove(labels_path)
except OSError:
pass
if __name__ == '__main__':