jittor-Anarcho/augment.py

import argparse
import cv2
import os
from tqdm import tqdm

def flip(root: str, file_name: str):
    img = cv2.imread(os.path.join(root, file_name))
    flipped = cv2.flip(img, 1)  # 水平翻转
    cv2.imwrite(os.path.join(root, "flipped_"+file_name), flipped)

def label_flip(root: str, file_name: str):
    img = cv2.imread(os.path.join(root, file_name), cv2.CV_16UC1)
    flipped = cv2.flip(img, 1)
    cv2.imwrite(os.path.join(root, "flipped_"+file_name), flipped)

def check_head(path: str):
    if path.startswith("flipped_"):
        return False
    return True

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_path', type=str)
    parser.add_argument("--process", type=int)
    args = parser.parse_args()

    # image
    for root, dirs, files in os.walk(args.input_path + "/imgs"):
        for file in tqdm(files):
            if check_head(file):
                flip(root, file)
                # print(os.path.join(root, file))

    # labels
    for root, dirs, files in tqdm(os.walk(args.input_path + "/labels")):
        for file in tqdm(files):
            if check_head(file):
                label_flip(root, file)