source https github comjason9075opencvmosaicdataaug import glob import os import random from string import asciilowercase digits import cv2 import numpy as np parameters outputsize 720 1280 height width scalerange 0 4 0 6 if height or width lower than this scale drop it filtertinyscale 1 100 labeldir imgdir outputdir numberimages 250 def main none imgpaths annos getdatasetlabeldir imgdir for index in rangenumberimages idxs random samplerangelenannos 4 newimage newannos path updateimageandanno imgpaths annos idxs outputsize scalerange filterscalefiltertinyscale get random string code 7b7ad245cdff75241935e4dd860f3bad lettercode randomchars32 filename path splitos sep1 rsplit 10 fileroot foutputdirfilenamemosaiclettercode cv2 imwriteffileroot jpg newimage cv2 imwritejpegquality 85 printfsucceeded index1numberimages with filename annoslist for anno in newannos width anno3 anno1 height anno4 anno2 xcenter anno1 width 2 ycenter anno2 height 2 obj fanno0 xcenter ycenter width height annoslist appendobj with openffileroot txt w as outfile outfile writen joinline for line in annoslist def getdatasetlabeldir str imgdir str tuplelist list imgpaths labels for labelfile in glob globos path joinlabeldir txt labelname labelfile splitos sep1 rsplit 10 with openlabelfile as infile objlists infile readlines imgpath os path joinimgdir flabelname jpg boxes for objlist in objlists obj objlist rstripn split xmin floatobj1 floatobj3 2 ymin floatobj2 floatobj4 2 xmax floatobj1 floatobj3 2 ymax floatobj2 floatobj4 2 boxes appendintobj0 xmin ymin xmax ymax if not boxes continue imgpaths appendimgpath labels appendboxes return imgpaths labels def updateimageandanno allimglist list allannos list idxs listint outputsize tupleint int scalerange tuplefloat float filterscale float 0 0 tuplelist list str outputimg np zerosoutputsize0 outputsize1 3 dtypenp uint8 scalex scalerange0 random random scalerange1 scalerange0 scaley scalerange0 random random scalerange1 scalerange0 dividpointx intscalex outputsize1 dividpointy intscaley outputsize0 newanno pathlist for i index in enumerateidxs path allimglistindex pathlist appendpath imgannos allannosindex img cv2 imreadpath if i 0 topleft img cv2 resizeimg dividpointx dividpointy outputimg dividpointy dividpointx img for bbox in imgannos xmin bbox1 scalex ymin bbox2 scaley xmax bbox3 scalex ymax bbox4 scaley newanno appendbbox0 xmin ymin xmax ymax elif i 1 topright img cv2 resizeimg outputsize1 dividpointx dividpointy outputimg dividpointy dividpointx outputsize1 img for bbox in imgannos xmin scalex bbox1 1 scalex ymin bbox2 scaley xmax scalex bbox3 1 scalex ymax bbox4 scaley newanno appendbbox0 xmin ymin xmax ymax elif i 2 bottomleft img cv2 resizeimg dividpointx outputsize0 dividpointy outputimgdividpointy outputsize0 dividpointx img for bbox in imgannos xmin bbox1 scalex ymin scaley bbox2 1 scaley xmax bbox3 scalex ymax scaley bbox4 1 scaley newanno appendbbox0 xmin ymin xmax ymax else bottomright img cv2 resize img outputsize1 dividpointx outputsize0 dividpointy outputimg dividpointy outputsize0 dividpointx outputsize1 img for bbox in imgannos xmin scalex bbox1 1 scalex ymin scaley bbox2 1 scaley xmax scalex bbox3 1 scalex ymax scaley bbox4 1 scaley newanno appendbbox0 xmin ymin xmax ymax remove bounding box small than scale of filter if filterscale 0 newanno anno for anno in newanno if filterscale anno3 anno1 and filterscale anno4 anno2 return outputimg newanno pathlist0 def randomcharsnumberchar int str assert numberchar 1 the number of character should greater than 1 lettercode asciilowercase digits return joinrandom choicelettercode for in rangenumberchar if name main main printdone parameters height width if height or width lower than this scale drop it get images list and annotations list from input dir update new images and annotations save images and annotations in output dir get random string code 7b7ad245cdff75241935e4dd860f3bad 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 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 top left top right bottom left bottom right remove bounding box small than scale of filter automatic generate random 32 characters get random string code 7b7ad245cdff75241935e4dd860f3bad len random_chars 32 32
import glob
import os
import random
from string import ascii_lowercase, digits

import cv2
import numpy as np


OUTPUT_SIZE = (720, 1280)  
SCALE_RANGE = (0.4, 0.6)  
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,
        )

        
        letter_code = random_chars(32)
        file_name = path.split(os.sep)[-1].rsplit(".", 1)[0]
        file_root = f"{OUTPUT_DIR}/{file_name}_MOSAIC_{letter_code}"
        cv2.imwrite(f"{file_root}.jpg", new_image, [cv2.IMWRITE_JPEG_QUALITY, 85])
        print(f"Succeeded {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(f"{file_root}.txt", "w") as outfile:
            outfile.write("\n".join(line for line in annos_list))


def get_dataset(label_dir: str, img_dir: str) -> tuple[list, list]:
    
    img_paths = []
    labels = []
    for label_file in glob.glob(os.path.join(label_dir, "*.txt")):
        label_name = label_file.split(os.sep)[-1].rsplit(".", 1)[0]
        with open(label_file) as in_file:
            obj_lists = in_file.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: list[int],
    output_size: tuple[int, int],
    scale_range: tuple[float, float],
    filter_scale: float = 0.0,
) -> tuple[list, list, str]:
    
    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, index in enumerate(idxs):
        path = all_img_list[index]
        path_list.append(path)
        img_annos = all_annos[index]
        img = cv2.imread(path)
        if i == 0:  
            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:  
            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:  
            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:  
            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])

    
    if filter_scale > 0:
        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: int) -> str:
    
    assert number_char > 1, "The number of character should greater than 1"
    letter_code = ascii_lowercase + digits
    return "".join(random.choice(letter_code) for _ in range(number_char))


if __name__ == "__main__":
    main()
    print("DONE ✅")