mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 11:46:35 +00:00
Oleksandr Bezdieniezhnykh
108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
import os.path
|
|
import time
|
|
from array import *
|
|
from pathlib import Path
|
|
import albumentations as A
|
|
import cv2
|
|
|
|
|
|
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')
|
|
|
|
|
|
class ImageLabel:
|
|
|
|
def __init__(self, image_path, image, labels_path, labels):
|
|
self.image_path = image_path
|
|
self.image = image
|
|
self.labels_path = labels_path
|
|
self.labels = labels
|
|
|
|
|
|
def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
|
|
transforms = [
|
|
A.Compose([A.HorizontalFlip(always_apply=True)], bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.RandomBrightnessContrast(always_apply=True)], bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=-40, always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=-20, always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=20, always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=40, always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo'))
|
|
]
|
|
|
|
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}'
|
|
results.append(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')
|
|
))
|
|
return results
|
|
|
|
|
|
def write_result(img_ann: ImageLabel):
|
|
cv2.imwrite(img_ann.image_path, img_ann.image)
|
|
with open(img_ann.labels_path, 'w') as f:
|
|
lines = [f'{ann[4]} {round(ann[0], 5)} {round(ann[1], 5)} {round(ann[2], 5)} {round(ann[3], 5)}\n' for ann in img_ann.labels]
|
|
f.writelines(lines)
|
|
f.close()
|
|
|
|
|
|
def read_labels(labels_path) -> [[]]:
|
|
with open(labels_path, 'r') as f:
|
|
rows = f.readlines()
|
|
arr = []
|
|
for row in rows:
|
|
str_coordinates = row.split(' ')
|
|
class_num = str_coordinates.pop(0)
|
|
coordinates = [float(n) for n in str_coordinates]
|
|
coordinates.append(class_num)
|
|
arr.append(coordinates)
|
|
return arr
|
|
|
|
|
|
def process_image(img_ann):
|
|
results = image_processing(img_ann)
|
|
for res_ann in results:
|
|
write_result(res_ann)
|
|
|
|
write_result(ImageLabel(
|
|
image=img_ann.image,
|
|
labels=img_ann.labels,
|
|
image_path=os.path.join(current_images_dir, Path(img_ann.image_path).name),
|
|
labels_path=os.path.join(current_labels_dir, Path(img_ann.labels_path).name)
|
|
))
|
|
os.remove(img_ann.image_path)
|
|
os.remove(img_ann.labels_path)
|
|
|
|
|
|
def main():
|
|
while True:
|
|
images = os.listdir(images_dir)
|
|
if len(images) == 0:
|
|
time.sleep(5)
|
|
continue
|
|
|
|
for image in images:
|
|
image_path = os.path.join(images_dir, image)
|
|
labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt')
|
|
process_image(ImageLabel(
|
|
image_path=image_path,
|
|
image=cv2.imread(image_path),
|
|
labels_path=labels_path,
|
|
labels=read_labels(labels_path)
|
|
))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |