ADD the algorithms of image augmentation

computer_vision/flip_augmentation.py

@@ -0,0 +1,109 @@
+import random
+import cv2
+import os
+import glob
+from typing import List
+
+"""
+Flip image and bounding box for computer vision task
+https://paperswithcode.com/method/randomhorizontalflip
+"""
+
+# Params
+LABEL_DIR = ''
+IMAGE_DIR = ''
+OUTPUT_DIR = ''
+FLIP_TYPE = 1  # (0 is vertical, 1 is horizontal)
+
+
+def main() -> None:
+    img_paths, annos = get_dataset(LABEL_DIR, IMAGE_DIR)
+    print('Processing...')
+    new_image, new_annos, path = update_image_and_anno(
+        img_paths, annos, FLIP_TYPE)
+
+    for index in range(len(new_image)):
+        # Get random string code: '7b7ad245cdff75241935e4dd860f3bad'
+        letter_code = random_chars(32)
+        file_name = path[index].split('/')[-1].rsplit('.', 1)[0]
+        cv2.imwrite(OUTPUT_DIR + f"/{file_name}_FLIP_{letter_code}.jpg", new_image[index], [cv2.IMWRITE_JPEG_QUALITY, 85])
+        print(f'Success {index+1}/{len(new_image)} with {file_name}')
+        annos_list = []
+        for anno in new_annos[index]:
+            obj = f'{anno[0]} {anno[1]} {anno[2]} {anno[3]} {anno[4]}'
+            annos_list.append(obj)
+        with open(OUTPUT_DIR + f"/{file_name}_FLIP{letter_code}.txt", "w") as outfile:
+            outfile.write("\n".join(line for line in annos_list))
+
+
+def get_dataset(label_dir: str, img_dir: str) -> List:
+    """
+    - label_dir <type: str>: Path to label include annotation of images
+    - img_dir <type: str>: Path to folder contain images
+    Return <type: list>: List of images path and labels 
+    """
+    img_paths = []
+    labels = []
+    for label_file in glob.glob(os.path.join(label_dir, '*.txt')):
+        label_name = label_file.split('/')[-1].rsplit('.', 1)[0]
+        f = open(label_file, 'r')
+        obj_lists = f.readlines()
+        img_path = os.path.join(img_dir, f'{label_name}.jpg')
+
+        boxes = []
+        for obj_list in obj_lists:
+            obj = obj_list.rstrip('\n').split(' ')
+            boxes.append([int(obj[0]), float(obj[1]),
+                          float(obj[2]), float(obj[3]), float(obj[4])])
+        if not boxes:
+            continue
+        img_paths.append(img_path)
+        labels.append(boxes)
+    return img_paths, labels
+
+
+def update_image_and_anno(img_list: List, anno_list: List, flip_type: int=1) -> List:
+    """
+    - img_list <type: list>: list of all images
+    - anno_list <type: list>: list of all annotations of specific image
+    - flip_type <type: int>: 0 is vertical, 1 is horizontal
+    Return:
+        - new_imgs_list <type: narray>: image after resize
+        - new_annos_lists <type: list>: list of new annotation after scale
+        - path_list <type: list>: list the name of image file  
+    """
+    new_annos_lists = []
+    path_list = []
+    new_imgs_list = []
+    for idx in range(len(img_list)):
+        new_annos = []
+        path = img_list[idx]
+        path_list.append(path)
+        img_annos = anno_list[idx]
+        img = cv2.imread(path)
+        if flip_type == 1:
+            new_img = cv2.flip(img, flip_type)
+            for bbox in img_annos:
+                x_center_new = 1 - bbox[1]
+                new_annos.append(
+                    [bbox[0], x_center_new, bbox[2], bbox[3], bbox[4]])
+        elif flip_type == 0:
+            new_img = cv2.flip(img, flip_type)
+            for bbox in img_annos:
+                y_center_new = 1 - bbox[2]
+                new_annos.append(
+                    [bbox[0], bbox[1], y_center_new, bbox[3], bbox[4]])
+        new_annos_lists.append(new_annos)
+        new_imgs_list.append(new_img)
+    return new_imgs_list, new_annos_lists, path_list
+
+
+def random_chars(number_char):
+    # Get random string code: '7b7ad245cdff75241935e4dd860f3bad'
+    letter_code = 'abcdefghijklmnopqrstuvwxyz0123456789'
+    return ''.join(random.choice(letter_code) for _ in range(number_char))
+
+
+if __name__ == '__main__':
+    main()
+    print('DONE ✅')

computer_vision/mosaic_augmentation.py

@@ -0,0 +1,161 @@
+'''Source: https://github.com/jason9075/opencv-mosaic-data-aug'''
+
+import random
+import cv2
+import os
+import glob
+import numpy as np
+from typing import List
+
+# Parrameters
+OUTPUT_SIZE = (720, 1280)  # Height, Width
+SCALE_RANGE = (0.4, 0.6)  # if height or width lower than this scale, drop it.
+FILTER_TINY_SCALE = 1 / 100
+LABEL_DIR = ''
+IMG_DIR = ''
+OUTPUT_DIR = ''
+NUMBER_IMAGES = 250
+
+
+def main() -> None:
+    img_paths, annos = get_dataset(LABEL_DIR, IMG_DIR)
+    for index in range(NUMBER_IMAGES):
+        idxs = random.sample(range(len(annos)), 4)
+        new_image, new_annos, path = update_image_and_anno(img_paths, annos,
+                                                           idxs,
+                                                           OUTPUT_SIZE, SCALE_RANGE,
+                                                           filter_scale=FILTER_TINY_SCALE)
+
+        # Get random string code: '7b7ad245cdff75241935e4dd860f3bad'
+        letter_code = random_chars(32)
+        file_name = path.split('/')[-1].rsplit('.', 1)[0]
+        cv2.imwrite(OUTPUT_DIR + f"/{file_name}_MOSAIC_{letter_code}.jpg", new_image, [cv2.IMWRITE_JPEG_QUALITY, 85])
+        print(f'Successed {index+1}/{NUMBER_IMAGES} with {file_name}')
+        annos_list = []
+        for anno in new_annos:
+            width = anno[3] - anno[1]
+            height = anno[4] - anno[2]
+            x_center = anno[1] + width/2
+            y_center = anno[2] + height/2
+            obj = f'{anno[0]} {x_center} {y_center} {width} {height}'
+            annos_list.append(obj)
+        with open(OUTPUT_DIR + f"/{file_name}_MOSAIC_{letter_code}.txt", "w") as outfile:
+            outfile.write("\n".join(line for line in annos_list))
+
+
+def get_dataset(label_dir: str, img_dir: str) -> List:
+    """
+    - label_dir <type: str>: Path to label include annotation of images
+    - img_dir <type: str>: Path to folder contain images
+    Return <type: list>: List of images path and labels
+    """
+    img_paths = []
+    labels = []
+    for label_file in glob.glob(os.path.join(label_dir, '*.txt')):
+        label_name = label_file.split('/')[-1].rsplit('.', 1)[0]
+        f = open(label_file, 'r')
+        obj_lists = f.readlines()
+        img_path = os.path.join(img_dir, f'{label_name}.jpg')
+
+        boxes = []
+        for obj_list in obj_lists:
+            obj = obj_list.rstrip('\n').split(' ')
+            xmin = float(obj[1]) - float(obj[3])/2
+            ymin = float(obj[2]) - float(obj[4])/2
+            xmax = float(obj[1]) + float(obj[3])/2
+            ymax = float(obj[2]) + float(obj[4])/2
+
+            boxes.append([int(obj[0]), xmin, ymin, xmax, ymax])
+        if not boxes:
+            continue
+        img_paths.append(img_path)
+        labels.append(boxes)
+    return img_paths, labels
+
+
+def update_image_and_anno(all_img_list: List, all_annos: List, idxs: int, output_size: int, scale_range: int, filter_scale: int=0.) -> List:
+    """
+    - all_img_list <type: list>: list of all images
+    - all_annos <type: list>: list of all annotations of specific image
+    - idxs <type: list>: index of image in list
+    - output_size <type: tuple>: size of output image (Height, Width)
+    - scale_range <type: tuple>: range of scale image
+    - filter_scale <type: float>: the condition of downscale image and bounding box
+    Return:
+        - output_img <type: narray>: image after resize
+        - new_anno <type: list>: list of new annotation after scale
+        - path[0] <type: string>: get the name of image file
+    """
+    output_img = np.zeros([output_size[0], output_size[1], 3], dtype=np.uint8)
+    scale_x = scale_range[0] + \
+        random.random() * (scale_range[1] - scale_range[0])
+    scale_y = scale_range[0] + \
+        random.random() * (scale_range[1] - scale_range[0])
+    divid_point_x = int(scale_x * output_size[1])
+    divid_point_y = int(scale_y * output_size[0])
+
+    new_anno = []
+    path_list = []
+    for i, idx in enumerate(idxs):
+        path = all_img_list[idx]
+        path_list.append(path)
+        img_annos = all_annos[idx]
+        img = cv2.imread(path)
+        if i == 0:  # top-left
+            img = cv2.resize(img, (divid_point_x, divid_point_y))
+            output_img[:divid_point_y, :divid_point_x, :] = img
+            for bbox in img_annos:
+                xmin = bbox[1] * scale_x
+                ymin = bbox[2] * scale_y
+                xmax = bbox[3] * scale_x
+                ymax = bbox[4] * scale_y
+                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])
+        elif i == 1:  # top-right
+            img = cv2.resize(
+                img, (output_size[1] - divid_point_x, divid_point_y))
+            output_img[:divid_point_y, divid_point_x:output_size[1], :] = img
+            for bbox in img_annos:
+                xmin = scale_x + bbox[1] * (1 - scale_x)
+                ymin = bbox[2] * scale_y
+                xmax = scale_x + bbox[3] * (1 - scale_x)
+                ymax = bbox[4] * scale_y
+                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])
+        elif i == 2:  # bottom-left
+            img = cv2.resize(
+                img, (divid_point_x, output_size[0] - divid_point_y))
+            output_img[divid_point_y:output_size[0], :divid_point_x, :] = img
+            for bbox in img_annos:
+                xmin = bbox[1] * scale_x
+                ymin = scale_y + bbox[2] * (1 - scale_y)
+                xmax = bbox[3] * scale_x
+                ymax = scale_y + bbox[4] * (1 - scale_y)
+                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])
+        else:  # bottom-right
+            img = cv2.resize(
+                img, (output_size[1] - divid_point_x, output_size[0] - divid_point_y))
+            output_img[divid_point_y:output_size[0],
+                       divid_point_x:output_size[1], :] = img
+            for bbox in img_annos:
+                xmin = scale_x + bbox[1] * (1 - scale_x)
+                ymin = scale_y + bbox[2] * (1 - scale_y)
+                xmax = scale_x + bbox[3] * (1 - scale_x)
+                ymax = scale_y + bbox[4] * (1 - scale_y)
+                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])
+
+    # Remove bounding box small than scale of filter
+    if 0 < filter_scale:
+        new_anno = [anno for anno in new_anno if
+                    filter_scale < (anno[3] - anno[1]) and filter_scale < (anno[4] - anno[2])]
+
+    return output_img, new_anno, path_list[0]
+
+
+def random_chars(number_char):
+    # Get random string code: '7b7ad245cdff75241935e4dd860f3bad'
+    letter_code = 'abcdefghijklmnopqrstuvwxyz0123456789'
+    return ''.join(random.choice(letter_code) for _ in range(number_char))
+
+
+if __name__ == '__main__':
+    main()
+    print('DONE ✅')