import os import time import math import random import csv from io import BytesIO import numpy as np from cairosvg import svg2png import cv2 import filetype from filetype.match import image_matchers from progress.bar import ChargingBar import imgaug as ia from imgaug import augmenters as iaa from imgaug.augmentables.batches import UnnormalizedBatch from entity import Entity from common import defaults, mkdir import imtool import pipelines BATCH_SIZE = 16 PARALLEL = 20 MIN_BACKGROUND_SIZE = 500 def process_bg(b): imw = cv2.imread(b.path) im, bb = imtool.remove_white(imw) annot = None label = b.path.replace('png', 'txt') if os.path.exists(label): # rewrite label with new coordinates [ww, wh, _] = imw.shape [iw, ih, _] = im.shape es = imtool.read_centroids(label) l = '' for e in es: [i, p, c] = e.values() [x,y,w,h] = [ max((c.x*ww - bb.x)/iw, 0), max((c.y*wh - bb.y)/ih, 0), (c.w*ww)/iw, (c.h*wh)/ih ] l += f'{int(i)} {x} {y} {w} {h}\n' annot = l if im.shape[0] > args.minbgsize and im.shape[1]> args.minbgsize: return im, annot else: raise Exception(f'droping {b.path} after remove_white => {im.shape}') def filter_bgs(bgs): ret = [] for b in bgs: if b.path.endswith('txt'): continue try: img, annot = process_bg(b) except Exception as e: print(f'drop: {e}') continue ret.append((b, img, annot)) return ret def process(args): dest_images_path = os.path.join(args.dest, 'images') dest_labels_path = os.path.join(args.dest, 'labels') mkdir.make_dirs([dest_images_path, dest_labels_path]) logo_images = [] logo_alphas = [] logo_labels = {} db = {} with open(defaults.MAIN_CSV_PATH, 'r') as f: reader = csv.DictReader(f) db = {e.bco: e for e in [Entity.from_dict(d) for d in reader]} background_images = [] for d in args.background: background_images.extend(os.scandir(d)) print(f'filtering {len(background_images)} background images from {args.background}') background_images = filter_bgs(background_images) assert(len(background_images)) stats = { 'failed': 0, 'ok': 0 } for d in os.scandir(args.logos): img = None if not d.is_file(): stats['failed'] += 1 continue try: if filetype.match(d.path, matchers=image_matchers): img = cv2.imread(d.path, cv2.IMREAD_UNCHANGED) else: png = svg2png(url=d.path) img = cv2.imdecode(np.asarray(bytearray(png), dtype=np.uint8), cv2.IMREAD_UNCHANGED) label = db[d.name.split('.')[0]].id (h, w, c) = img.shape if c == 3: img = imtool.add_alpha(img) if img.ndim < 3: print(f'very bad dim: {img.ndim}') (h, w, c) = img.shape assert(w > 10) assert(h > 10) stats['ok'] += 1 (b, g, r, _) = cv2.split(img) alpha = img[:, :, 3]/255 d = cv2.merge([b, g, r]) logo_images.append(d) # tried id() tried __array_interface__, tried tagging, nothing works logo_labels.update({d.tobytes(): label}) # XXX(xaiki): we pass alpha as a float32 heatmap, # because imgaug is pretty strict about what data it will process # and that we want the alpha layer to pass the same transformations as the orig logo_alphas.append(np.dstack((alpha, alpha, alpha)).astype('float32')) except Exception as e: stats['failed'] += 1 print(f'error loading: {d.path}: {e}') print(stats) assert(len(logo_alphas) == len(logo_images)) print(f"will process {len(logo_images)} images on {len(background_images)} backgrounds") # so that we don't get a lot of the same logos on the same page. zipped = list(zip(logo_images, logo_alphas)) random.shuffle(zipped) logo_images, logo_alphas = zip(*zipped) n = len(logo_images) batches = [] for i in range(math.floor(n*2/BATCH_SIZE)): s = (i*BATCH_SIZE)%n e = min(s + BATCH_SIZE, n) le = max(0, BATCH_SIZE - (e - s)) a = logo_images[0:le] + logo_images[s:e] h = logo_alphas[0:le] + logo_alphas[s:e] assert(len(a) == BATCH_SIZE) batches.append(UnnormalizedBatch(images=a,heatmaps=h)) bar = ChargingBar(f'augment ({len(logo_images)} logos {len(background_images)} bgs)', max=(len(batches)**2)/3*len(background_images)) # We use a single, very fast augmenter here to show that batches # are only loaded once there is space again in the buffer. pipeline = pipelines.HUGE def create_generator(lst): for b in lst: print(f"Loading next unaugmented batch...") yield b batches_generator = create_generator(batches) batch = 0 with pipeline.pool(processes=args.parallel, seed=1) as pool: batches_aug = pool.imap_batches(batches_generator, output_buffer_size=5) print(f"Requesting next augmented batch...{batch}/{len(batches)}") for i, batch_aug in enumerate(batches_aug): idx = list(range(len(batch_aug.images_aug))) random.shuffle(idx) for j, (d, img, annot) in enumerate(background_images): basename = d.name.replace('.png', f'.{i}.{j}') annotations = [] try: annotations.append(annot.rstrip()) except: pass for k in range(math.floor(len(batch_aug.images_aug)/3)): bar.next() logo_idx = (j+k*4)%len(batch_aug.images_aug) orig = batch_aug.images_unaug[logo_idx] label = logo_labels[orig.tobytes()] logo = batch_aug.images_aug[logo_idx] assert(logo.shape == orig.shape) # XXX(xaiki): we get alpha from heatmap, but will only use one channel # we could make mix_alpha into mix_mask and pass all 3 chanels alpha = cv2.split(batch_aug.heatmaps_aug[logo_idx]) try: bb = imtool.mix_alpha(img, logo, alpha[0], random.random(), random.random()) c = bb.to_centroid(img.shape) annotations.append(c.to_annotation(label)) except AssertionError as err: print(f'couldnt process {i}, {j}: {err}') except Exception as err: print(f'error in mix pipeline: {err}') try: cv2.imwrite(f'{dest_images_path}/{basename}.png', img) label_path = f"{dest_labels_path}/{basename}.txt" with open(label_path, 'a') as f: f.write('\n'.join(annotations)) except Exception: print(f'couldnt write image {basename}') if i < len(batches)-1: print(f"Requesting next augmented batch...{batch}/{len(batches)}") batch += 1 bar.finish() if __name__ == '__main__': import argparse print("✨ augmenting data") parser = argparse.ArgumentParser(description='mix backgrounds and logos into augmented data for YOLO') parser.add_argument('--logos', metavar='logos', type=str, default=defaults.LOGOS_DATA_PATH, help='dir containing logos') parser.add_argument('--background', metavar='backgrounds', type=str, nargs='+', default=[defaults.SCREENSHOT_PATH, defaults.FISH_PATH], help='dir containing background plates') parser.add_argument('--dst', dest='dest', type=str, default=defaults.AUGMENTED_DATA_PATH, help='dest dir') parser.add_argument('--parallel', metavar='parallel', type=int, default=PARALLEL, help='number of concurrent jobs') parser.add_argument('--min-background-size', dest='minbgsize', type=int, default=MIN_BACKGROUND_SIZE, help='minimum background size') args = parser.parse_args() process(args)