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 check.py

Browse Source
This commit is contained in:
Nazar Sturanec
2024-05-25 03:13:19 +03:00
parent 8529bd0c6f
commit 6adc4c37e9
21 changed files with 95 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
preprocessing.py
+23 -19
View File
@@ -2,13 +2,12 @@ import os.path
import albumentations as A
import cv2
from pathlib import Path
import datetime
labels_dir = 'labels'
images_dir = 'images'
current_dataset_dir = os.path.join('datasets', 'zombobase-current')
class ImageAnnotation:
def read_annotations(self) -> [[]]:
with open(self.dataset_annotation_path, 'r') as f:
rows = f.readlines()
arr = []
@@ -21,7 +20,7 @@ class ImageAnnotation:
return arr
def __init__(self, image_path):
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')
@@ -32,9 +31,9 @@ class ImageAnnotation:
self.annotations = self.read_annotations()
def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
def image_processing(img_ann, current_dataset_dir,labels_dir,images_dir: ImageAnnotation) -> [ImageAnnotation]:
category_ids = []
bboxes = ImageAnnotation(img_ann).read_annotations()
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)
@@ -45,7 +44,7 @@ def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids']))
bboxes = bboxes
imag = ImageAnnotation(img_ann).image
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']
@@ -54,18 +53,24 @@ def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
return transformed_image, transformed_bboxes, transformed_category_ids
def write_results(img_ann: ImageAnnotation):
for i in range(100):
transformed_image, transformed_bboxes, transformed_category_ids, = image_processing(img_ann)
cv2.imwrite(os.path.join(current_dataset_dir, images_dir, str(i)+ImageAnnotation(img_ann).image_path + '.jpg'), transformed_image)
with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann).image_path + '.txt'), 'w') as f:
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():
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)):
@@ -80,10 +85,9 @@ def process_image():
return file_annotation_finished
def main():
process_image()
for i in process_image():
write_results(i)
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)
if __name__ == '__main__':
main()
#main(os.path.join('Zombobase-' + str(datetime.date.today()), 'datasets', 'zombobase-current','test'),'labels', 'images')