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

refactor augmentation to class, update classes.json, fix small bugs

Browse Source
This commit is contained in:
zxsanny
2025-05-15 06:10:10 +03:00
parent fd4cd75265
commit 292809593c
5 changed files with 188 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
augmentation.py
+160
View File
@@ -0,0 +1,160 @@
import concurrent.futures
import os.path
import shutil
import time
from datetime import datetime
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, processed_dir)
from dto.imageLabel import ImageLabel
class Augmentator:
def __init__(self):
self.total_files_processed = 0
self.total_images_to_process = 0
self.correct_margin = 0.0005
self.correct_min_bbox_size = 0.01
self.transform = A.Compose([
# Flips, rotations and brightness
A.HorizontalFlip(p=0.6),
A.RandomBrightnessContrast(p=0.4, brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1)),
A.Affine(p=0.7, scale=(0.8, 1.2), rotate=(-20, 20), shear=(-10, 10), translate_percent=0.2),
# Weather
A.RandomFog(p=0.3, fog_coef_range=(0, 0.3)),
A.RandomShadow(p=0.2),
# Image Quality/Noise
A.MotionBlur(p=0.2, blur_limit=(3, 5)),
# Color Variations
A.HueSaturationValue(p=0.4, hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10)
], bbox_params=A.BboxParams(format='yolo'))
def correct_bboxes(self, labels):
res = []
for bboxes in labels:
x = bboxes[0]
y = bboxes[1]
half_width = 0.5*bboxes[2]
half_height = 0.5*bboxes[3]
# calc how much bboxes are outside borders ( +small margin ).
# value should be negative. If it's positive, then put 0, as no correction
w_diff = min((1 - self.correct_margin) - (x + half_width), (x - half_width) - self.correct_margin, 0)
w = bboxes[2] + 2*w_diff
if w < self.correct_min_bbox_size:
continue
h_diff = min((1 - self.correct_margin) - (y + half_height), ((y - half_height) - self.correct_margin), 0)
h = bboxes[3] + 2 * h_diff
if h < self.correct_min_bbox_size:
continue
res.append([x, y, w, h, bboxes[4]])
return res
pass
def augment_inner(self, img_ann: ImageLabel) -> [ImageLabel]:
results = []
labels = self.correct_bboxes(img_ann.labels)
if len(labels) == 0 and len(img_ann.labels) != 0:
print('no labels but was!!!')
results.append(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)
)
)
for i in range(7):
try:
res = self.transform(image=img_ann.image, bboxes=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 read_labels(self, 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 augment_annotation(self, image_file):
try:
image_path = os.path.join(data_images_dir, image_file.name)
labels_path = os.path.join(data_labels_dir, f'{Path(str(image_path)).stem}.txt')
image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
img_ann = ImageLabel(
image_path=image_path,
image=image,
labels_path=labels_path,
labels=self.read_labels(labels_path)
)
try:
results = self.augment_inner(img_ann)
for annotation in results:
cv2.imencode('.jpg', annotation.image)[1].tofile(annotation.image_path)
with open(annotation.labels_path, 'w') as f:
lines = [f'{l[4]} {round(l[0], 5)} {round(l[1], 5)} {round(l[2], 5)} {round(l[3], 5)}\n' for l in
annotation.labels]
f.writelines(lines)
f.close()
print(f'{datetime.now():{'%Y-%m-%d %H:%M:%S'}}: {self.total_files_processed + 1}/{self.total_images_to_process} : {image_file.name} has augmented')
except Exception as e:
print(e)
self.total_files_processed += 1
except Exception as e:
print(f'Error appeared in thread for {image_file.name}: {e}')
def augment_annotations(self, from_scratch=False):
self.total_files_processed = 0
if from_scratch:
shutil.rmtree(processed_dir)
os.makedirs(processed_images_dir, exist_ok=True)
os.makedirs(processed_labels_dir, exist_ok=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)
self.total_images_to_process = len(images)
with concurrent.futures.ThreadPoolExecutor() as executor:
executor.map(self.augment_annotation, images)
if __name__ == '__main__':
augmentator = Augmentator()
while True:
augmentator.augment_annotations()
print('All processed, waiting for 2 minutes...')
time.sleep(300)
classes.json
+16 -16
View File
@@ -1,18 +1,18 @@
[
{ "Id": 0, "Name": "Armored-Vehicle", "Color": "#80FF0000" },
{ "Id": 1, "Name": "Truck", "Color": "#8000FF00" },
{ "Id": 2, "Name": "Vehicle", "Color": "#800000FF" },
{ "Id": 3, "Name": "Artillery", "Color": "#80FFFF00" },
{ "Id": 4, "Name": "Shadow", "Color": "#80FF00FF" },
{ "Id": 5, "Name": "Trenches", "Color": "#8000FFFF" },
{ "Id": 6, "Name": "Military-men", "Color": "#80188021" },
{ "Id": 7, "Name": "Tyre-tracks", "Color": "#80800000" },
{ "Id": 8, "Name": "Additional-armored-tank", "Color": "#80008000" },
{ "Id": 9, "Name": "Smoke", "Color": "#80000080" },
{ "Id": 10, "Name": "Plane", "Color": "#80000080" },
{ "Id": 11, "Name": "Moto", "Color": "#80808000" },
{ "Id": 12, "Name": "Camouflage-net", "Color": "#80800080" },
{ "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" },
{ "Id": 14, "Name": "Roof", "Color": "#80008080" },
{ "Id": 15, "Name": "Building", "Color": "#80008080" }
{ "Id": 0, "Name": "ArmorVehicle", "ShortName": "Броня", "Color": "#FF0000" },
{ "Id": 1, "Name": "Truck", "ShortName": "Вантаж.", "Color": "#00FF00" },
{ "Id": 2, "Name": "Vehicle", "ShortName": "Машина", "Color": "#0000FF" },
{ "Id": 3, "Name": "Atillery", "ShortName": "Арта", "Color": "#FFFF00" },
{ "Id": 4, "Name": "Shadow", "ShortName": "Тінь", "Color": "#FF00FF" },
{ "Id": 5, "Name": "Trenches", "ShortName": "Окопи", "Color": "#00FFFF" },
{ "Id": 6, "Name": "MilitaryMan", "ShortName": "Військов", "Color": "#188021" },
{ "Id": 7, "Name": "TyreTracks", "ShortName": "Накати", "Color": "#800000" },
{ "Id": 8, "Name": "AdditArmoredTank", "ShortName": "Танк.захист", "Color": "#008000" },
{ "Id": 9, "Name": "Smoke", "ShortName": "Дим", "Color": "#000080" },
{ "Id": 10, "Name": "Plane", "ShortName": "Літак", "Color": "#000080" },
{ "Id": 11, "Name": "Moto", "ShortName": "Мото", "Color": "#808000" },
{ "Id": 12, "Name": "CamouflageNet", "ShortName": "Сітка", "Color": "#800080" },
{ "Id": 13, "Name": "CamouflageBranches", "ShortName": "Гілки", "Color": "#2f4f4f" },
{ "Id": 14, "Name": "Roof", "ShortName": "Дах", "Color": "#1e90ff" },
{ "Id": 15, "Name": "Building", "ShortName": "Будівля", "Color": "#ffb6c1" }
]
preprocess-train.py
+3 -4
View File
@@ -1,6 +1,5 @@
from preprocessing import preprocess_annotations
from augmentation import Augmentator
from train import train_dataset, convert2rknn
preprocess_annotations()
train_dataset('2024-10-01')
convert2rknn()
Augmentator().augment_annotations(from_scratch=True)
train_dataset(from_scratch=True)
preprocessing.py
-162
View File
@@ -1,162 +0,0 @@
import concurrent.futures
import os.path
import time
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)
from dto.imageLabel import ImageLabel
total_files_processed = 0
transform = A.Compose([
# Flips, rotations and brightness
A.HorizontalFlip(),
A.RandomBrightnessContrast(brightness_limit=(-0.05, 0.05), contrast_limit=(-0.05, 0.05)),
A.Affine(p=0.7, scale=(0.8, 1.2), rotate=25, translate_percent=0.1),
# Weather
A.RandomFog(p=0.2, fog_coef_range=(0, 0.3)),
A.RandomShadow(p=0.2),
# Image Quality/Noise
A.MotionBlur(p=0.2, blur_limit=(3, 5)),
# Color Variations
A.HueSaturationValue(p=0.3, hue_shift_limit=8, sat_shift_limit=8, val_shift_limit=8)
], bbox_params=A.BboxParams(format='yolo'))
def correct_bboxes(labels):
margin = 0.0005
min_size = 0.01
res = []
for bboxes in labels:
x = bboxes[0]
y = bboxes[1]
half_width = 0.5*bboxes[2]
half_height = 0.5*bboxes[3]
# calc how much bboxes are outside borders ( +small margin ).
# value should be negative. If it's positive, then put 0, as no correction
w_diff = min((1 - margin) - (x + half_width), (x - half_width) - margin, 0)
w = bboxes[2] + 2*w_diff
if w < min_size:
continue
h_diff = min((1 - margin) - (y + half_height), ((y - half_height) - margin), 0)
h = bboxes[3] + 2 * h_diff
if h < min_size:
continue
res.append([x, y, w, h, bboxes[4]])
return res
pass
def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
results = []
labels = correct_bboxes(img_ann.labels)
if len(labels) == 0 and len(img_ann.labels) != 0:
print('no labels but was!!!')
results.append(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)
)
)
for i in range(7):
try:
res = transform(image=img_ann.image, bboxes=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):
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()
global total_files_processed
print(f'{total_files_processed}. {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 preprocess_annotations():
global total_files_processed # Indicate that we're using the global counter
total_files_processed = 0
os.makedirs(processed_images_dir, exist_ok=True)
os.makedirs(processed_labels_dir, exist_ok=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)
with concurrent.futures.ThreadPoolExecutor() as executor:
executor.map(process_image_file, images)
def process_image_file(image_file):
try:
image_path = os.path.join(data_images_dir, image_file.name)
labels_path = os.path.join(data_labels_dir, f'{Path(str(image_path)).stem}.txt')
image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
img_ann = ImageLabel(
image_path=image_path,
image=image,
labels_path=labels_path,
labels=read_labels(labels_path)
)
try:
results = image_processing(img_ann)
for res_ann in results:
write_result(res_ann)
except Exception as e:
print(e)
global total_files_processed
total_files_processed += 1
except Exception as e:
print(f'Error appeared in thread for {image_file.name}: {e}')
def main():
while True:
preprocess_annotations()
print('All processed, waiting for 2 minutes...')
time.sleep(300)
if __name__ == '__main__':
main()
train.py
+7 -6
View File
@@ -38,6 +38,7 @@ DEFAULT_CLASS_NUM = 80
total_files_copied = 0
def form_dataset(from_date: datetime):
makedirs(today_dataset, exist_ok=True)
images = []
old_images = []
@@ -67,7 +68,6 @@ def form_dataset(from_date: datetime):
copy_annotations(images[:train_size], 'train')
copy_annotations(images[train_size:train_size + valid_size], 'valid')
copy_annotations(images[train_size + valid_size:], 'test')
create_yaml()
def copy_annotations(images, folder):
@@ -174,21 +174,22 @@ def get_latest_model():
def train_dataset(existing_date=None, from_scratch=False):
latest_date, latest_model = get_latest_model()
latest_date, latest_model = get_latest_model() if not from_scratch else None, None
if existing_date is not None:
cur_folder = f'{prefix}{existing_date}'
cur_dataset = path.join(datasets_dir, f'{prefix}{existing_date}')
else:
# form_dataset(latest_date)
# create_yaml()
if from_scratch:
shutil.rmtree(today_dataset)
form_dataset(latest_date)
create_yaml()
cur_folder = today_folder
cur_dataset = today_dataset
model_name = latest_model if latest_model is not None and path.isfile(latest_model) and not from_scratch else 'yolo11m.yaml'
print(f'Initial model: {model_name}')
model = YOLO(model_name)
model.info['author'] = 'LLC Azaion'
yaml = abspath(path.join(cur_dataset, 'data.yaml'))
results = model.train(data=yaml,
@@ -222,7 +223,7 @@ def validate(model_path):
if __name__ == '__main__':
# model_path = train_dataset(from_scratch=True)
model_path = train_dataset(from_scratch=True)
# validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
# form_data_sample(500)
# convert2rknn()