mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-23 01:46:36 +00:00
add structure how it should work
in preprocessing.py
This commit is contained in:
Binary file not shown.
|
After Width: | Height: | Size: 2.3 MiB |
@@ -0,0 +1,3 @@
|
|||||||
|
3 0.41877 0.64332 0.06107 0.07926
|
||||||
|
7 0.40442 0.78827 0.06779 0.05212
|
||||||
|
7 0.19860 0.90662 0.05314 0.05429
|
||||||
@@ -0,0 +1,72 @@
|
|||||||
|
import os.path
|
||||||
|
import time
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import albumentations as alb
|
||||||
|
from os import listdir
|
||||||
|
from os.path import isfile, join
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
labels_dir = 'labels'
|
||||||
|
images_dir = 'images'
|
||||||
|
current_dataset_dir = os.path.join('datasets', 'zombobase-current')
|
||||||
|
|
||||||
|
class ImageAnnotation:
|
||||||
|
|
||||||
|
def read_annotations(self) -> [[]]:
|
||||||
|
with open(self.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):
|
||||||
|
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_name, os.path.basename(image_path))
|
||||||
|
self.image = cv2.imread(image_path)
|
||||||
|
|
||||||
|
self.annotation_path = os.path.join(labels_dir, self.image_name, '.txt')
|
||||||
|
self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_name, '.txt')
|
||||||
|
self.annotations = self.read_annotations()
|
||||||
|
|
||||||
|
def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
|
||||||
|
# return structure example:
|
||||||
|
# utilize transform albumentations here
|
||||||
|
return [ImageAnnotation(f'{img_ann.image_name}1', image1, bboxes1 ),
|
||||||
|
ImageAnnotation(f'{img_ann.image_name}2', image2, bboxes2),
|
||||||
|
...
|
||||||
|
]
|
||||||
|
|
||||||
|
def write_results(img_ann: ImageAnnotation):
|
||||||
|
# write image cv2.imwrite(, image) dataset_image_path
|
||||||
|
# write img_ann.annotations into new file with name dataset_annotation_path
|
||||||
|
|
||||||
|
|
||||||
|
def process_image(img_ann):
|
||||||
|
results = image_processing(img_ann)
|
||||||
|
for res_ann in results:
|
||||||
|
write_results(res_ann)
|
||||||
|
write_results(img_ann)
|
||||||
|
os.remove(img_ann.image_path)
|
||||||
|
os.remove(img_ann.annotation_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)
|
||||||
|
process_image(ImageAnnotation(image_path))
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user