mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 11:46:35 +00:00
Nazar Sturanec
93 lines
4.0 KiB
Python
93 lines
4.0 KiB
Python
import os.path
|
|
import albumentations as A
|
|
import cv2
|
|
from pathlib import Path
|
|
import datetime
|
|
|
|
class ImageAnnotation:
|
|
|
|
def read_annotations(self) -> [[]]:
|
|
|
|
with open(self.dataset_annotation_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 __init__(self, image_path, current_dataset_dir, labels_dir, images_dir):
|
|
self.image_path = image_path
|
|
self.image_name = Path(image_path).stem
|
|
self.dataset_image_path = os.path.join(current_dataset_dir, images_dir, self.image_path + '.jpg')
|
|
self.image = cv2.imread(self.dataset_image_path)
|
|
|
|
self.annotation_path = os.path.join(labels_dir, self.image_path + '.txt')
|
|
self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_path + '.txt')
|
|
self.annotations = self.read_annotations()
|
|
|
|
|
|
def image_processing(img_ann, current_dataset_dir,labels_dir,images_dir: ImageAnnotation) -> [ImageAnnotation]:
|
|
category_ids = []
|
|
bboxes = ImageAnnotation(img_ann,current_dataset_dir,labels_dir,images_dir).read_annotations()
|
|
for i in range(len(bboxes)):
|
|
category_ids.append(bboxes[i][4])
|
|
bboxes[i].pop(4)
|
|
|
|
transform = A.Compose([
|
|
A.HorizontalFlip(p=0.5),
|
|
A.RandomBrightnessContrast(p=0.2),
|
|
], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids']))
|
|
|
|
bboxes = bboxes
|
|
imag = ImageAnnotation(img_ann,current_dataset_dir,labels_dir,images_dir).image
|
|
|
|
transformed = transform(image=imag, bboxes=bboxes, category_ids=category_ids)
|
|
transformed_image = transformed['image']
|
|
transformed_bboxes = transformed['bboxes']
|
|
transformed_category_ids = transformed['category_ids']
|
|
|
|
return transformed_image, transformed_bboxes, transformed_category_ids
|
|
|
|
|
|
def write_results(img_ann, current_dataset_dir, labels_dir, images_dir: ImageAnnotation):
|
|
|
|
file_start_save = 'Zombobase-' + str(datetime.date.today())
|
|
for i in range(5):
|
|
|
|
transformed_image, transformed_bboxes, transformed_category_ids, = image_processing(img_ann, current_dataset_dir, labels_dir, images_dir)
|
|
cv2.imwrite(os.path.join(current_dataset_dir, images_dir,str(i)+ImageAnnotation(img_ann,current_dataset_dir,labels_dir,images_dir).image_path + '.jpg'), transformed_image)
|
|
with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann, current_dataset_dir, labels_dir, images_dir).image_path + '.txt'), 'w') as f:
|
|
print(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann, current_dataset_dir, labels_dir, images_dir).image_path + '.txt'))
|
|
for bbox, category_id in zip(transformed_bboxes, transformed_category_ids):
|
|
print(bbox)
|
|
x_center, y_center, width, height = bbox
|
|
cla = category_id
|
|
f.write(f"{cla} {x_center} {y_center} {width} {height}\n")
|
|
#
|
|
#
|
|
def process_image(current_dataset_dir,images_dir):
|
|
file_annotation = []
|
|
file_annotation_finished =[]
|
|
for foldername, subfolders, filenames in os.walk(os.path.join(current_dataset_dir,images_dir)):
|
|
file_annotation.append(filenames)
|
|
print(file_annotation)
|
|
|
|
for i in range(len(file_annotation[0])):
|
|
ff = file_annotation[0][i].split('.')
|
|
ff.pop(-1)
|
|
ff = '.'.join(ff)
|
|
file_annotation_finished.append(ff)
|
|
|
|
return file_annotation_finished
|
|
|
|
def main(current_dataset_dir, labels_dir, images_dir):
|
|
process_image(current_dataset_dir,images_dir)
|
|
for i in process_image(current_dataset_dir,images_dir):
|
|
write_results(i, current_dataset_dir, labels_dir, images_dir)
|
|
|
|
#main(os.path.join('Zombobase-' + str(datetime.date.today()), 'datasets', 'zombobase-current','test'),'labels', 'images') |