diff --git a/AnnotationClass.py b/dto/annotationClass.py similarity index 76% rename from AnnotationClass.py rename to dto/annotationClass.py index b04ff58..ef8938c 100644 --- a/AnnotationClass.py +++ b/dto/annotationClass.py @@ -1,4 +1,5 @@ import json +from os.path import dirname, join class AnnotationClass: @@ -9,7 +10,8 @@ class AnnotationClass: @staticmethod def read_json(): - with open('classes.json', 'r', encoding='utf-8') as f: + classes_path = join(dirname(dirname(__file__)), 'classes.json') + with open(classes_path, 'r', encoding='utf-8') as f: j = json.loads(f.read()) return {cl['Id']: AnnotationClass(id=cl['Id'], name=cl['Name'], color=cl['Color']) for cl in j} diff --git a/dto/imageLabel.py b/dto/imageLabel.py new file mode 100644 index 0000000..798052f --- /dev/null +++ b/dto/imageLabel.py @@ -0,0 +1,32 @@ +import cv2 +from matplotlib import pyplot as plt + + +class ImageLabel: + def __init__(self, image_path, image, labels_path, labels): + self.image_path = image_path + self.image = image + self.labels_path = labels_path + self.labels = labels + + def visualize(self, annotation_classes): + img = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2RGB) + height, width, channels = img.shape + for label in self.labels: + class_num = int(label[-1]) + x_c = float(label[0]) + y_c = float(label[1]) + w = float(label[2]) + h = float(label[3]) + x_min = x_c - w / 2 + y_min = y_c - h / 2 + x_max = x_min + w + y_max = y_min + h + color = annotation_classes[class_num].color_tuple + + cv2.rectangle(img, (int(x_min * width), int(y_min * height)), (int(x_max * width), int(y_max * height)), + color=color, thickness=3) + plt.figure(figsize=(12, 12)) + plt.axis('off') + plt.imshow(img) + plt.show() \ No newline at end of file diff --git a/models/zombobase_2024-05-18.pt b/models/zombobase_2024-05-18.pt deleted file mode 100644 index d61ef50..0000000 Binary files a/models/zombobase_2024-05-18.pt and /dev/null differ diff --git a/preprocessing.py b/preprocessing.py index 795e0e5..f476a57 100644 --- a/preprocessing.py +++ b/preprocessing.py @@ -3,8 +3,8 @@ import time from pathlib import Path import albumentations as A import cv2 -from matplotlib import pyplot as plt -from AnnotationClass import AnnotationClass +from dto.annotationClass import AnnotationClass +from dto.imageLabel import ImageLabel labels_dir = 'labels' images_dir = 'images' @@ -14,36 +14,6 @@ current_labels_dir = os.path.join(current_dataset_dir, 'labels') annotation_classes = AnnotationClass.read_json() -class ImageLabel: - def __init__(self, image_path, image, labels_path, labels): - self.image_path = image_path - self.image = image - self.labels_path = labels_path - self.labels = labels - - def visualize(self): - img = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2RGB) - height, width, channels = img.shape - for label in self.labels: - class_num = int(label[-1]) - x_c = float(label[0]) - y_c = float(label[1]) - w = float(label[2]) - h = float(label[3]) - x_min = x_c - w / 2 - y_min = y_c - h / 2 - x_max = x_min + w - y_max = y_min + h - color = annotation_classes[class_num].color_tuple - - cv2.rectangle(img, (int(x_min * width), int(y_min * height)), (int(x_max * width), int(y_max * height)), - color=color, thickness=3) - plt.figure(figsize=(12, 12)) - plt.axis('off') - plt.imshow(img) - plt.show() - - def image_processing(img_ann: ImageLabel) -> [ImageLabel]: transforms = [ A.Compose([A.HorizontalFlip(always_apply=True)], @@ -85,13 +55,16 @@ def write_result(img_ann: ImageLabel, show_image=False): os.makedirs(os.path.dirname(img_ann.labels_path), exist_ok=True) if show_image: - img_ann.visualize() + img_ann.visualize(annotation_classes) cv2.imwrite(img_ann.image_path, img_ann.image) + 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() + print(f'{img_ann.labels_path} written') def read_labels(labels_path) -> [[]]: @@ -129,14 +102,17 @@ def main(): continue for image in images: - image_path = os.path.join(images_dir, image) - labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt') - process_image(ImageLabel( - image_path=image_path, - image=cv2.imread(image_path), - labels_path=labels_path, - labels=read_labels(labels_path) - )) + try: + image_path = os.path.join(images_dir, image) + labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt') + process_image(ImageLabel( + image_path=image_path, + image=cv2.imread(image_path), + labels_path=labels_path, + labels=read_labels(labels_path) + )) + except FileNotFoundError: + print(f'No labels file {labels_path} found') if __name__ == '__main__': diff --git a/tests/imagelabel_visualize_test.py b/tests/imagelabel_visualize_test.py new file mode 100644 index 0000000..67c7890 --- /dev/null +++ b/tests/imagelabel_visualize_test.py @@ -0,0 +1,25 @@ +from pathlib import Path +import cv2 +import os.path + +from dto.annotationClass import AnnotationClass +from dto.imageLabel import ImageLabel +from preprocessing import read_labels + +images_dir = '../images' +labels_dir = '../labels' +annotation_classes = AnnotationClass.read_json() + + +image = os.listdir(images_dir)[0] +image_path = os.path.join(images_dir, image) +labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt') + +img = ImageLabel( + image_path=image_path, + image=cv2.imread(image_path), + labels_path=labels_path, + labels=read_labels(labels_path) +) + +img.visualize(annotation_classes) diff --git a/train.py b/train.py new file mode 100644 index 0000000..8105dbd --- /dev/null +++ b/train.py @@ -0,0 +1,4 @@ +import os + + +current_dataset_dir = os.path.join('datasets', 'zombobase-current')