neuro-mkr/aug.py

import albumentations as a
import numpy as np
from os import listdir as ls
import cv2 as cv

def b_c(b, s = (1920, 1080)):
    return [
        round(b[0] * s[0] - b[2] * s[0] / 2),
        round(b[0] * s[0] + b[2] * s[0] / 2),
        round(b[1] * s[1] - b[3] * s[1] / 2),
        round(b[1] * s[1] + b[3] * s[1] / 2),
    ]

def sq(img, bbox):
    for i in bbox:
        b = b_c(i, s = (img.shape[1], img.shape[0]))
        #print(b)
        cv.rectangle(img, (b[0], b[2]), (b[1], b[3]), (255, 0, 0), 2)

base_path = '../videos/total'
tagt_path = '../videos/total-exp'

p = 1

t = a.Compose([
    a.SafeRotate(limit=80, p=0.8),
    a.BBoxSafeRandomCrop(),
    a.HorizontalFlip(p=0.5),
    a.VerticalFlip(p=0.5),
    a.RandomBrightnessContrast(p=0.9, brightness_limit = (-0.2, 0.7), contrast_limit = (-0.6, 0.4)),
    a.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.8),
    a.ChromaticAberration(p=0.2),
], bbox_params=a.BboxParams(format='yolo', label_fields=['class_labels']))

fs = set([i.rsplit(".", 1)[0] for i in ls(base_path) if i != "classes.txt"])

for i in list(fs):
    img = cv.cvtColor(cv.imread(f"{base_path}/{i}.jpg"), cv.COLOR_BGR2RGB)
    box = [tuple(map(float, j.split()[1:])) for j in open(f"{base_path}/{i}.txt").read().split("\n") if j]
    
    lbl = np.array(['ch'] * len(box))
    box = np.array(box)

    for k in range(10):
        r = t(image = img, bboxes = box, class_labels = lbl)

        out_img = r['image']
        out_box = r['bboxes']
        out_lbl = r['class_labels']

        print(f"open({tagt_path}/{p}-{i}-{k}.txt)")
        #print("\n".join([" ".join(list(map(str, ['0'] + list(y)))) for y in out_box]) + "\n")
        open(f"{tagt_path}/{p}-{i}-{k}.txt", 'w').write("\n".join([" ".join(list(map(str, ['0'] + list(y)))) for y in out_box]) + "\n")
        cv.imwrite(f"{tagt_path}/{p}-{i}-{k}.jpg", cv.cvtColor(out_img, cv.COLOR_RGB2BGR))

    #sq(out_img, out_box)
    
    #from matplotlib import pyplot as plt
    #plt.imshow(out_img)
    #plt.show()