File size: 3,596 Bytes
0e83290 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import cv2
import uuid
import shutil
import numpy as np
from pathlib import Path
from typing import Protocol, Type
from tools.contour_detector import getting_coordinates
from tools.mask_display import mask_map
SAVE_FOLDER = Path.cwd() / 'video-test'
class ExportObject:
def __init__(self, mask, name_class) -> None:
self.mask = mask
self.name_class = name_class
class ExportImage:
def __init__(self, image, objects: list[ExportObject]) -> None:
self.image = image
self.exports_objects = objects
class TypeSave(Protocol):
def start_creation(self) -> None:
pass
def create_archive(self) -> str:
pass
def get_type_save_annotation(
images: list[np.ndarray],
masks: list[np.ndarray],
names_class: list[str],
type_save: str = 'yolo',
) -> TypeSave:
'''Factory'''
types_saves: dict[str, Type[TypeSave]] = {
'yolo': YoloSave,
}
return types_saves[type_save](images, masks, names_class)
class YoloSave:
def __init__(
self, images: list[np.ndarray], masks: list[np.ndarray], names_class: list[str]
) -> None:
self.images = images
self.masks = masks
self.names_class = {}
for i, name in enumerate(names_class):
self.names_class[name] = i
folder_name = (
''.join(names_class)[:10]
if len(''.join(names_class)) > 15
else ''.join(names_class)
)
folder_name = f'dt-{folder_name}-{uuid.uuid1()}'
p = Path(SAVE_FOLDER / folder_name)
p.mkdir()
self.path_folder = SAVE_FOLDER / folder_name
self.images_folder = self.path_folder / 'images'
p = Path(self.images_folder)
p.mkdir()
self.lables_folder = self.path_folder / 'lables'
p = Path(self.lables_folder)
p.mkdir()
def start_creation(self):
path_image = self.images_folder / 'image_filename'
path_txt = self.lables_folder / 'image_filename'
for i, (image, mask) in enumerate(zip(self.images, self.masks)):
cv2.imwrite(f'{path_image}{i+1}.jpg', image)
txt_frame_save(image, mask, f'{path_txt}{i+1}', 0)
txt_class_save(self.path_folder / 'classes', self.names_class)
def create_archive(self) -> str:
shutil.make_archive(self.path_folder, 'zip', self.path_folder)
shutil.rmtree(self.path_folder)
return f'{self.path_folder}.zip'
def txt_class_save(path: str, names_class: dict):
with open(f'{path}.txt', 'w') as file:
for key, value in names_class.items():
name_class_str = [f'{value} {key} \n']
file.writelines(name_class_str)
def txt_frame_save(
images: np.ndarray, mask_unique: np.ndarray, path: str, name_class_idx: int
):
img_height = images.shape[0]
img_width = images.shape[1]
with open(f'{path}.txt', 'w') as file:
coordinates = []
for mask in mask_map(mask_unique):
coordinate = getting_coordinates(mask)
coordinates += coordinate
for box in coordinates:
x, y = box[0], box[1]
w, h = box[2], box[3]
x_center = x + int(w / 2)
y_center = y + int(h / 2)
norm_xc = x_center / img_width
norm_yc = y_center / img_height
norm_width = w / img_width
norm_height = h / img_height
yolo_annotation = [
f'{name_class_idx} {norm_xc} {norm_yc} {norm_width} {norm_height} \n'
]
file.writelines(yolo_annotation)
|