mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 12:06:36 +00:00
zxsanny
122 lines
4.6 KiB
Python
122 lines
4.6 KiB
Python
import os.path
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
import albumentations as A
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir,
|
|
annotation_classes, checkpoint_file, checkpoint_date_format)
|
|
from dto.imageLabel import ImageLabel
|
|
|
|
|
|
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.SafeRotate(limit=90, always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.SafeRotate(limit=90, always_apply=True),
|
|
A.RandomBrightnessContrast(always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True),
|
|
A.VerticalFlip(always_apply=True), ],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo')),
|
|
A.Compose([A.SafeRotate(limit=90, always_apply=True),
|
|
A.RandomBrightnessContrast(always_apply=True)],
|
|
bbox_params=A.BboxParams(format='yolo'))
|
|
]
|
|
|
|
results = []
|
|
for i, transform in enumerate(transforms):
|
|
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(processed_images_dir, f'{name}{path.suffix}'),
|
|
labels_path=os.path.join(processed_labels_dir, f'{name}.txt')
|
|
)
|
|
results.append(img)
|
|
except Exception as e:
|
|
print(f'Error during transformation: {e}')
|
|
return results
|
|
|
|
|
|
def write_result(img_ann: ImageLabel):
|
|
os.makedirs(os.path.dirname(img_ann.image_path), exist_ok=True)
|
|
os.makedirs(os.path.dirname(img_ann.labels_path), exist_ok=True)
|
|
|
|
cv2.imencode('.jpg', img_ann.image)[1].tofile(img_ann.image_path)
|
|
print(f'{img_ann.image_path} written')
|
|
|
|
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()
|
|
print(f'{img_ann.labels_path} written')
|
|
|
|
|
|
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.replace(',', '.')) for n in str_coordinates]
|
|
# noinspection PyTypeChecker
|
|
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(processed_images_dir, Path(img_ann.image_path).name),
|
|
labels_path=os.path.join(processed_labels_dir, Path(img_ann.labels_path).name)
|
|
))
|
|
|
|
|
|
def main():
|
|
while True:
|
|
processed_images = set(f.name for f in os.scandir(processed_images_dir))
|
|
images = []
|
|
with os.scandir(data_images_dir) as imd:
|
|
for image_file in imd:
|
|
if image_file.is_file() and image_file.name not in processed_images:
|
|
images.append(image_file)
|
|
|
|
for image_file in images:
|
|
try:
|
|
image_path = os.path.join(data_images_dir, image_file.name)
|
|
labels_path = os.path.join(data_labels_dir, f'{Path(image_path).stem}.txt')
|
|
image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
|
|
process_image(ImageLabel(
|
|
image_path=image_path,
|
|
image=image,
|
|
labels_path=labels_path,
|
|
labels=read_labels(labels_path)
|
|
))
|
|
except Exception as e:
|
|
print(f'Error appeared {e}')
|
|
print('All processed, waiting for 2 minutes...')
|
|
time.sleep(120)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|