├── .DS_Store ├── coco_classes ├── 15_cat.png ├── 16_dog.png ├── 2_car.png ├── 62_tv.png ├── 8_boat.png ├── 0_person.png ├── 22_zebra.png ├── 53_pizza.png ├── 54_donut.png ├── 6_train.png ├── 74_clock.png ├── 23_giraffe.png ├── 46_banana.png ├── 48_sandwich.png ├── 49_orange.png ├── 4_airplane.png ├── 61_toilet.png ├── 11_stop sign.png ├── 10_fire hydrant.png └── 72_refrigerator.png ├── utils ├── get_img_txt.py ├── coco2yolo.py └── yolov5_2_coco.py ├── classes.txt ├── remove_bg.py ├── class_info.py ├── uncertainly.py ├── baseline1_aug.py ├── README.md ├── select.py ├── baseline2_aug.py ├── img_add.py └── gen_mask_img.py /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/.DS_Store -------------------------------------------------------------------------------- /coco_classes/15_cat.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/15_cat.png -------------------------------------------------------------------------------- /coco_classes/16_dog.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/16_dog.png -------------------------------------------------------------------------------- /coco_classes/2_car.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/2_car.png -------------------------------------------------------------------------------- /coco_classes/62_tv.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/62_tv.png -------------------------------------------------------------------------------- /coco_classes/8_boat.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/8_boat.png -------------------------------------------------------------------------------- /coco_classes/0_person.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/0_person.png -------------------------------------------------------------------------------- /coco_classes/22_zebra.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/22_zebra.png -------------------------------------------------------------------------------- /coco_classes/53_pizza.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/53_pizza.png -------------------------------------------------------------------------------- /coco_classes/54_donut.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/54_donut.png -------------------------------------------------------------------------------- /coco_classes/6_train.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/6_train.png -------------------------------------------------------------------------------- /coco_classes/74_clock.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/74_clock.png -------------------------------------------------------------------------------- /coco_classes/23_giraffe.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/23_giraffe.png -------------------------------------------------------------------------------- /coco_classes/46_banana.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/46_banana.png -------------------------------------------------------------------------------- /coco_classes/48_sandwich.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/48_sandwich.png -------------------------------------------------------------------------------- /coco_classes/49_orange.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/49_orange.png -------------------------------------------------------------------------------- /coco_classes/4_airplane.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/4_airplane.png -------------------------------------------------------------------------------- /coco_classes/61_toilet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/61_toilet.png -------------------------------------------------------------------------------- /coco_classes/11_stop sign.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/11_stop sign.png -------------------------------------------------------------------------------- /coco_classes/10_fire hydrant.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/10_fire hydrant.png -------------------------------------------------------------------------------- /coco_classes/72_refrigerator.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSCT-Lab/ObjTest/HEAD/coco_classes/72_refrigerator.png -------------------------------------------------------------------------------- /utils/get_img_txt.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | 5 | if __name__ == '__main__': 6 | parser = argparse.ArgumentParser() 7 | parser.add_argument('--img_path', default='E:/code/datasets/coco1000/images/val2017', required=True) 8 | parser.add_argument('--root_path', default='E:/code/datasets/coco1000', required=True) 9 | parser.add_argument('--txt_name', default='coco1000_val2017.txt', required=True) 10 | args = parser.parse_args() 11 | 12 | save_li = [] 13 | dir_path = args.img_path.split(args.root_path)[1] 14 | path_li = os.listdir(args.img_path) 15 | for p in path_li: 16 | # save_li.append('./' + dir_path + '/' + p) 17 | save_li.append('.' + dir_path + '/' + p) 18 | 19 | save_path = os.path.join(args.root_path, args.txt_name) 20 | with open(save_path, 'w') as f: 21 | f.writelines([line + '\n' for line in save_li]) -------------------------------------------------------------------------------- /classes.txt: -------------------------------------------------------------------------------- 1 | 0: person 2 | 1: bicycle 3 | 2: car 4 | 3: motorcycle 5 | 4: airplane 6 | 5: bus 7 | 6: train 8 | 7: truck 9 | 8: boat 10 | 9: traffic light 11 | 10: fire hydrant 12 | 11: stop sign 13 | 12: parking meter 14 | 13: bench 15 | 14: bird 16 | 15: cat 17 | 16: dog 18 | 17: horse 19 | 18: sheep 20 | 19: cow 21 | 20: elephant 22 | 21: bear 23 | 22: zebra 24 | 23: giraffe 25 | 24: backpack 26 | 25: umbrella 27 | 26: handbag 28 | 27: tie 29 | 28: suitcase 30 | 29: frisbee 31 | 30: skis 32 | 31: snowboard 33 | 32: sports ball 34 | 33: kite 35 | 34: baseball bat 36 | 35: baseball glove 37 | 36: skateboard 38 | 37: surfboard 39 | 38: tennis racket 40 | 39: bottle 41 | 40: wine glass 42 | 41: cup 43 | 42: fork 44 | 43: knife 45 | 44: spoon 46 | 45: bowl 47 | 46: banana 48 | 47: apple 49 | 48: sandwich 50 | 49: orange 51 | 50: broccoli 52 | 51: carrot 53 | 52: hot dog 54 | 53: pizza 55 | 54: donut 56 | 55: cake 57 | 56: chair 58 | 57: couch 59 | 58: potted plant 60 | 59: bed 61 | 60: dining table 62 | 61: toilet 63 | 62: tv 64 | 63: laptop 65 | 64: mouse 66 | 65: remote 67 | 66: keyboard 68 | 67: cell phone 69 | 68: microwave 70 | 69: oven 71 | 70: toaster 72 | 71: sink 73 | 72: refrigerator 74 | 73: book 75 | 74: clock 76 | 75: vase 77 | 76: scissors 78 | 77: teddy bear 79 | 78: hair drier 80 | 79: toothbrush -------------------------------------------------------------------------------- /remove_bg.py: -------------------------------------------------------------------------------- 1 | from rembg import remove 2 | import os 3 | import argparse 4 | 5 | 6 | def remove_bg(input_dir, output_dir): 7 | for img_file in os.listdir(input_dir): 8 | fname = img_file.split('.jpg')[0] 9 | input_path = os.path.join(input_dir, img_file) 10 | output_path = os.path.join(output_dir, fname + '.png') # 保存透明背景的图片为 png 11 | 12 | with open(input_path, 'rb') as i: 13 | with open(output_path, 'wb') as o: 14 | input = i.read() 15 | output = remove(input) 16 | o.write(output) 17 | 18 | 19 | if __name__ == '__main__': 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument('--dataset', type=str, choices=['coco', 'openimages', 'voc'], required=True) 22 | arg = parser.parse_args() 23 | 24 | input_dir, output_dir = '', '' 25 | if arg.dataset == 'coco': 26 | input_dir = './coco_last500' # 原始图片的位置 27 | output_dir = './coco_last500_remove' # 去除背景后，透明背景底的图片存储位置 28 | elif arg.dataset == 'openimages': 29 | input_dir = './openimages_random500' 30 | output_dir = './openimages_random500_remove' 31 | elif arg.dataset == 'voc': 32 | input_dir = './voc_random500' 33 | output_dir = './voc_random500_remove' 34 | else: 35 | print("This dataset is not currently supported!") 36 | 37 | os.makedirs(output_dir, exist_ok=True) 38 | remove_bg(input_dir, output_dir) 39 | -------------------------------------------------------------------------------- /class_info.py: -------------------------------------------------------------------------------- 1 | coco_class_dict = { 2 | 0: '0_person.png', 3 | 2: '2_car.png', 4 | 4: '4_airplane.png', 5 | 6: '6_train.png', 6 | 8: '8_boat.png', 7 | 10: '10_fire hydrant.png', 8 | 11: '11_stop sign.png', 9 | 15: '15_cat.png', 10 | 16: '16_dog.png', 11 | 22: '22_zebra.png', 12 | 23: '23_giraffe.png', 13 | 46: '46_banana.png', 14 | 48: '48_sandwich.png', 15 | 49: '49_orange.png', 16 | 53: '53_pizza.png', 17 | 54: '54_donut.png', 18 | 61: '61_toilet.png', 19 | 62: '62_tv.png', 20 | 72: '72_refrigerator.png', 21 | 74: '74_clock.png' 22 | } 23 | 24 | openimages_class_dict = { 25 | 110: '110_Snack.png', 26 | 121: '121_cat.png', 27 | 228: '228_door handle.png', 28 | 34: '34_Sunglasses.png', 29 | 404: '404_handbag.png', 30 | 426: '426_Turtle.png', 31 | 433: '433_Footwear.png', 32 | 440: '440_wind glass.png', 33 | 445: '445_dog.png', 34 | 455: '455_Flower.png', 35 | 467: '467_Airplane.png', 36 | 483: '483_Moths and butterflies.png', 37 | 498: '498_Camera.png', 38 | 503: '503_Vegetable.png', 39 | 510: '510_Cabbage.png', 40 | 570: '570_car.png', 41 | 592: '592_Alpaca.png', 42 | 595: '595_Remote control.png', 43 | 68: '68_person.png', 44 | 98: '98_poster.png' 45 | } 46 | 47 | voc_class_dict = { 48 | 0: '0_aeroplane.png', 49 | 1: '1_bicycle.png', 50 | 2: '2_bird.png', 51 | 4: '4_bottle.png', 52 | 5: '5_bus.png', 53 | 6: '6_car.png', 54 | 7: '7_cat.png', 55 | 11: '11_dog.png', 56 | 14: '14_person.png', 57 | 18: '18_train.png' 58 | } 59 | -------------------------------------------------------------------------------- /utils/coco2yolo.py: -------------------------------------------------------------------------------- 1 | """ 2 | author: Wu 3 | https://github.com/Weifeng-Chen/DL_tools/issues/3 4 | 2021/1/24 5 | COCO 格式的数据集转化为 YOLO 格式的数据集，源代码采取遍历方式，太慢， 6 | 这里改进了一下时间复杂度，从O(nm)改为O(n+m)，但是牺牲了一些内存占用 7 | --json_path 输入的json文件路径 8 | --save_path 保存的文件夹名字，默认为当前目录下的labels。 9 | """ 10 | 11 | import os 12 | import json 13 | from tqdm import tqdm 14 | import argparse 15 | 16 | 17 | def convert(size, box): 18 | dw = 1. / (size[0]) 19 | dh = 1. / (size[1]) 20 | x = box[0] + box[2] / 2.0 21 | y = box[1] + box[3] / 2.0 22 | w = box[2] 23 | h = box[3] 24 | 25 | x = x * dw 26 | w = w * dw 27 | y = y * dh 28 | h = h * dh 29 | return (x, y, w, h) 30 | 31 | 32 | if __name__ == '__main__': 33 | parser = argparse.ArgumentParser() 34 | parser.add_argument('--json_path', type=str, help="input: coco format(json)", required=True) 35 | parser.add_argument('--save_path', type=str, help="specify where to save the output dir of labels", required=True) 36 | arg = parser.parse_args() 37 | 38 | json_file = arg.json_path # COCO Object Instance 类型的标注 39 | ana_txt_save_path = arg.save_path # 保存的路径 40 | 41 | data = json.load(open(json_file, 'r')) 42 | if not os.path.exists(ana_txt_save_path): 43 | os.makedirs(ana_txt_save_path) 44 | 45 | id_map = {} # coco数据集的id不连续！重新映射一下再输出！ 46 | for i, category in enumerate(data['categories']): 47 | id_map[category['id']] = i 48 | 49 | # 通过事先建表来降低时间复杂度 50 | max_id = 0 51 | for img in data['images']: 52 | max_id = max(max_id, img['id']) 53 | # 注意这里不能写作 [[]]*(max_id+1)，否则列表内的空列表共享地址 54 | img_ann_dict = [[] for i in range(max_id + 1)] 55 | for i, ann in enumerate(data['annotations']): 56 | print(ann['image_id']) 57 | img_ann_dict[ann['image_id']].append(i) 58 | 59 | for img in tqdm(data['images']): 60 | filename = img["file_name"] 61 | img_width = img["width"] 62 | img_height = img["height"] 63 | img_id = img["id"] 64 | head, tail = os.path.splitext(filename) 65 | if '/' in head: 66 | head = head.split('/')[-1] 67 | ana_txt_name = head + ".txt" # 对应的txt名字，与jpg一致 68 | # print(ana_txt_name) 69 | f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w') 70 | '''for ann in data['annotations']: 71 | if ann['image_id'] == img_id: 72 | box = convert((img_width, img_height), ann["bbox"]) 73 | f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))''' 74 | # 这里可以直接查表而无需重复遍历 75 | for ann_id in img_ann_dict[img_id]: 76 | ann = data['annotations'][ann_id] 77 | box = convert((img_width, img_height), ann["bbox"]) 78 | f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3])) 79 | f_txt.close() -------------------------------------------------------------------------------- /uncertainly.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import shutil 4 | 5 | def label_correlation_sorted(file_path, adj): 6 | pic_info={} 7 | for filename in os.listdir(file_path): 8 | with open(file_path+"\\"+filename, 'r') as f: 9 | label_path=file_path+"\\"+filename 10 | print("label_path:") 11 | print(label_path) 12 | lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 13 | max_conf_label=lb[-1][0] 14 | adj_lb=adj[int(max_conf_label)] 15 | min_cor=1 16 | for x in lb: 17 | min_cor=min(adj_lb[int(x[0])],min_cor) 18 | pic_info[filename]=min_cor 19 | return(pic_info) 20 | 21 | 22 | def coffidence_sorted(file_path): 23 | pic_info = {} 24 | for filename in os.listdir(file_path): 25 | with open(file_path+"\\"+filename, 'r') as f: 26 | sum=0 27 | label_path=file_path+"\\"+filename 28 | print("label_path:") 29 | print(label_path) 30 | lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 31 | for x in lb: 32 | sum+=x[-1] 33 | avg_conf=sum/len(lb) 34 | pic_info[filename]=avg_conf 35 | return pic_info 36 | 37 | 38 | 39 | 40 | if __name__ == '__main__': 41 | # adj = np.load("./adj.npy") 42 | # label_path = r"C:\Users\Xu Jiali\Desktop\南软\测试拓展\ODdata\labels" 43 | # data_path = r"D:\dataset\coco\aug\labels" 44 | # dict = label_correlation_sorted(label_path , adj) 45 | # dict_1 = sorted(dict.items(),key=lambda x: x[1]) 46 | # print(dict_1) 47 | # with open("./sorted/coco/5l/label_corr.txt","w") as f: 48 | # for x in os.listdir(data_path): 49 | # if x not in os.listdir(label_path): 50 | # f.write("./aug/images/" + x.split('.')[0] + ".png" + '\n') 51 | # for x in dict_1: 52 | # f.write("./aug/images/" + str((x[0].split('.'))[0]) + ".png" +'\n') 53 | # f.close() 54 | 55 | 56 | # dict=coffidence_sorted(label_path) 57 | # dict_2 = sorted(dict.items(),key=lambda x: x[1]) 58 | # print(dict_2) 59 | # with open("./sorted/coco/5m_avg_conf.txt", "w") as f: 60 | # for x in os.listdir(data_path): 61 | # if x not in os.listdir(label_path): 62 | # f.write("./aug/images/" + x.split('.')[0] + ".png" + '\n') 63 | # for x in dict_2: 64 | # f.write("./aug/images/" + str((x[0].split('.'))[0]) + ".png" +'\n') 65 | # f.close() 66 | 67 | 68 | 69 | with open(r"./sorted/bdd/5m_avg_conf.txt", "r") as f: 70 | lines=f.readlines() 71 | path=r"D:\dataset\bdd\aug" 72 | i=0 73 | while i <= 600: # 150, 450, 600 74 | filename=lines[i].split('/')[3] 75 | filename=filename.split('.')[0] 76 | 77 | image_path=path+"\images\\"+filename+".png" 78 | image_new_path = r".\600\bdd\images" 79 | os.makedirs(image_new_path, exist_ok=True) 80 | label_path=path+"\labels\\"+filename+".txt" 81 | label_new_path=r".\600\bdd\labels" 82 | os.makedirs(label_new_path, exist_ok=True) 83 | shutil.copy(image_path, image_new_path) 84 | shutil.copy(label_path, label_new_path) 85 | i+=1 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | -------------------------------------------------------------------------------- /baseline1_aug.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import random 4 | import os 5 | import shutil 6 | import argparse 7 | 8 | 9 | # 图片亮度变换 10 | def brightness_aug(img_dir, save_img_dir): 11 | for img_file in os.listdir(img_dir): 12 | fname = img_file.split('.jpg')[0] 13 | source_path = os.path.join(img_dir, img_file) 14 | img = cv2.imread(source_path) 15 | 16 | parameter = random.choice(list(range(30, 80)) + list(range(-80, -30))) 17 | 18 | # 增加图像亮度 19 | res = np.uint8(np.clip((1 * np.int16(img) + parameter), 0, 255)) 20 | img_save_path = os.path.join(save_img_dir, fname + '.png') 21 | cv2.imwrite(img_save_path, res) # 随机改变图片亮度的新图片存入图片文件夹 22 | 23 | 24 | # 图片对比度变换 25 | def contrast_aug(img_dir, save_img_dir): 26 | for img_file in os.listdir(img_dir): 27 | fname = img_file.split('.jpg')[0] 28 | source_path = os.path.join(img_dir, img_file) 29 | 30 | parameter = random.uniform(1, 3) 31 | parameter = round(parameter, 1) 32 | 33 | img = cv2.imread(source_path) 34 | res = np.uint8(np.clip((parameter * (np.int16(img) - 60) + 50), 0, 255)) 35 | 36 | img_save_path = os.path.join(save_img_dir, fname + '.png') 37 | cv2.imwrite(img_save_path, res) # 随机改变图片亮度的新图片存入图片文件夹 38 | 39 | 40 | # 图片模糊度变换 41 | def blur_aug(img_dir, save_img_dir): 42 | for img_file in os.listdir(img_dir): 43 | fname = img_file.split('.jpg')[0] 44 | source_path = os.path.join(img_dir, img_file) 45 | 46 | kernel_size = (5, 5) 47 | sigma = random.uniform(2, 5) 48 | sigma = round(sigma, 1) 49 | 50 | img = cv2.imread(source_path) 51 | res = cv2.GaussianBlur(img, kernel_size, sigma) 52 | 53 | img_save_path = os.path.join(save_img_dir, fname + '.png') 54 | cv2.imwrite(img_save_path, res) # 随机改变图片模糊度的新图片存入图片文件夹 55 | 56 | 57 | if __name__ == '__main__': 58 | parser = argparse.ArgumentParser() 59 | parser.add_argument('--img_dir', type=str, default='../datasets/coco1000/images/val2017', help='img path') 60 | parser.add_argument('--label_dir', type=str, default='../datasets/coco1000/labels/val2017', help='label path') 61 | parser.add_argument('--save_dir', type=str, default='../datasets/obj_aug_coco1000', help='save path') 62 | arg = parser.parse_args() 63 | 64 | img_dir, label_dir, save_dir = arg.img_dir, arg.label_dir, arg.save_dir 65 | brightness_save_img_dir = os.path.join(save_dir, 'brightness/images') 66 | brightness_save_label_dir = os.path.join(save_dir, 'brightness/labels') 67 | contrast_save_img_dir = os.path.join(save_dir, 'contrast/images') 68 | contrast_save_label_dir = os.path.join(save_dir, 'contrast/labels') 69 | blur_save_img_dir = os.path.join(save_dir, 'blur/images') 70 | blur_save_label_dir = os.path.join(save_dir, 'blur/labels') 71 | 72 | os.makedirs(brightness_save_img_dir, exist_ok=True) 73 | os.makedirs(contrast_save_img_dir, exist_ok=True) 74 | os.makedirs(blur_save_img_dir, exist_ok=True) 75 | 76 | brightness_aug(img_dir, brightness_save_img_dir) 77 | contrast_aug(img_dir, contrast_save_img_dir) 78 | blur_aug(img_dir, blur_save_img_dir) 79 | 80 | shutil.copytree(label_dir, brightness_save_label_dir) 81 | shutil.copytree(label_dir, contrast_save_label_dir) 82 | shutil.copytree(label_dir, blur_save_label_dir) 83 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ObjTest: Object-Level Mutation for Testing Object Detection Systems 2 | 3 | ObjTest has conducted experiments on 3 datasets (COCO, VOC, BDD100K) and the testing subject is YOLOv5 with 3 different model sizes (YOLOv5s, YOLOv5m, YOLOv5l). 4 | 5 | ## Environment 6 | 7 | - python=3.9 8 | - Rembg 9 | 10 | We recommend using conda to configure the environment: 11 | 12 | ```sh 13 | conda create -n python=3.9 14 | conda activate objtest 15 | pip install rembg 16 | ``` 17 | 18 | ## Object-level Transformation 19 | 20 | We have implemented three object-level transformations: object insertion, removing, replacing. 21 | 22 | ### Object Insertion 23 | 24 | - **Step1**: Extract objects and remove their background: 25 | 26 | ```sh 27 | python remove_bg.py --dataset coco 28 | ``` 29 | 30 | Manually selecte some clear and clean images and save to ``ObjTest/coco_classes`` folder. In our demo repository, the object images have already saved in this folder. 31 | 32 | - **Step2**: Run ``img_add.py`` to overlay the mask image onto the original image and save the new image and the corresponding label: 33 | 34 | ````sh 35 | python img_add.py --task insertion --dataset coco --type val 36 | ```` 37 | 38 | The generated datasets after executing the insertion operator are in the ``E:/code/datasets/obj_aug_coco1000/insertion/images`` and ``E:/code/datasets/obj_aug_coco1000/insertion/ labels`` folders, they can be input to YOLOv5 directly. You can change the file path in the code yourself. 39 | 40 | ### Object Removing 41 | 42 | - **Step1**: Run ``ObjTest/gen_mask_img.py`` to generate the mask image corresponding to COCO1000 and the label after removing: 43 | 44 | ````sh 45 | python gen_mask_img.py --task remove --dataset coco --type val 46 | ```` 47 | 48 | - **Step2**: Download and configure the [lama](Download and configure the lama project) project. Then, go to the ``lama-main`` project and perform the fix for mask. You can change the floaer path by yourself. 49 | 50 | ````sh 51 | cd lama-main 52 | conda activate lama 53 | # coco val: 54 | python ./bin/predict.py model.path=E:/code/lama-main/big-lama indir=E:/code/datasets/obj_aug_coco1000/remove/images_tmp outdir=E:/code/datasets/obj_aug_coco1000/remove/images 55 | ```` 56 | 57 | The generated datasets after executing the remove operator are in the ``E:/code/datasets/obj_aug_coco1000/remove/images`` and ``E:/code/datasets/obj_aug_coco1000/remove/labels`` folders, they can be input to YOLOv5 directly. 58 | 59 | ### Object Replacing 60 | 61 | - **Step1**: prepare the mask image: 62 | 63 | ````sh 64 | python gen_mask_img.py --task replace --dataset coco --type val 65 | ```` 66 | 67 | After running, ``datasets/obj_aug_coco1000/replace/images_tmp1`` contains the original map and the corresponding mask location map; ``datasets/obj_aug_coco1000/replace/labels_tmp`` contains the new label after removing object from the original image and the deleted object's label file. 68 | 69 | - **Step2**: Use ``lama`` to repair the removed object image. You need to change the path location yourself. 70 | 71 | ````sh 72 | cd lama-main 73 | conda activate lama 74 | # coco val: 75 | python ./bin/predict.py model.path=E:/code/lama-main/big-lama indir=E:/code/datasets/obj_aug_coco1000/replace/images_tmp1 outdir=E:/code/datasets/obj_aug_coco1000/replace/images_tmp2 76 | ```` 77 | 78 | 最后，用 img_add.py增加 Obj，需要在``datasets/obj_aug_coco1000/replace/labels_tmp2``的基础上增加Obj，结合``datasets/obj_aug_coco1000/replace/labels_tmp/xxx_replace.txt``提供的类别和位置信息： 79 | 80 | - **Step3**: Add object with ``img_add.py``. It requires adding Object to ``datasets/obj_aug_coco1000/replace/labels_tmp2``, combined with ``datasets/obj_aug_coco1000/replace/labels_tmp /xxx_replace.txt`` to provide the category and location information. 81 | 82 | ````sh 83 | python img_add.py --task replace --dataset coco --type val 84 | ```` 85 | 86 | The generated datasets after executing the replace operator are in the ``E:/code/datasets/obj_aug_coco1000/replace/images`` and ``E:/code/datasets/obj_aug_coco1000/replace/labels`` folders, they can be input to YOLOv5 directly. 87 | 88 | ## Verification in YOLOv5 OD System 89 | 90 | In the YOLOv5 project, verify each of the generated dataset with three transformations: 91 | 92 | ````sh 93 | conda activate yolov5 94 | python val.py --weights ./torch_models/yolov5s.pt --data coco_insertion.yaml --img 640 --batch-size 4 --verbose 95 | python val.py --weights ./torch_models/yolov5s.pt --data coco_remove.yaml --img 640 --batch-size 4 --verbose 96 | python val.py --weights ./torch_models/yolov5s.pt --data coco_replace.yaml --img 640 --batch-size 4 --verbose 97 | ```` 98 | 99 | ## Baselines 100 | 101 | - **Baseline1**: The first type of baseline, which directly performs random transformations of brightness, contrast, and Gaussian blur on the image as a whole, without changing the label information. 102 | 103 | ````sh 104 | python baseline1_aug.py --img_dir ../datasets/coco1000/images/val2017 --label_dir ../datasets/coco1000/labels/val2017 --save_dir ../datasets/obj_aug_coco1000 105 | ```` 106 | 107 | - **Baseline2**: The second type of baseline, cutting and stitching of images, requires changing the label information. 108 | 109 | ````sh 110 | python baseline2_aug.py --img_dir ../datasets/coco1000/images/val2017 --label_dir ../datasets/coco1000/labels/val2017 --save_dir ../datasets/obj_aug_coco1000 111 | ```` 112 | 113 | ## Naturalness verification of generated images 114 | 115 | In the ``pytorch-fid`` folder, the FID score is calculated by inputing in the path of two image folders; the smaller the FID score, the more realistic and natural the image is. 116 | 117 | ````sh 118 | cd pytorch_fid 119 | conda activate fid 120 | run_coco.sh 121 | ```` 122 | 123 | -------------------------------------------------------------------------------- /select.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import shutil 4 | import numpy as np 5 | 6 | 7 | def random_sample(img_in_path, label_in_path, img_out_path, label_out_path, m=100): 8 | source_paths, txt_paths, fnames = [], [], [] 9 | for img_file in os.listdir(img_in_path): 10 | fname = img_file.split('.')[0] 11 | if os.path.exists(os.path.join(label_in_path, fname + '.txt')): # 找到匹配的 jpg 和 txt 12 | source_paths.append(os.path.join(img_in_path, img_file)) 13 | txt_paths.append(os.path.join(label_in_path, fname + '.txt')) 14 | fnames.append(fname) 15 | idx = list(range(0, len(source_paths))) 16 | selected_idx = random.sample(idx, m) 17 | for i in selected_idx: 18 | shutil.copyfile(source_paths[i], os.path.join(img_out_path, fnames[i] + '.jpg')) 19 | shutil.copyfile(txt_paths[i], os.path.join(label_out_path, fnames[i] + '.txt')) 20 | 21 | 22 | def get_aug(insertion_dir, remove_dir, replace_dir, out_dir): 23 | out_img_dir = os.path.join(out_dir, 'images') 24 | out_label_dir = os.path.join(out_dir, 'labels') 25 | os.makedirs(out_img_dir, exist_ok=True) 26 | os.makedirs(out_label_dir, exist_ok=True) 27 | 28 | insertion_img = os.path.join(insertion_dir, 'images') 29 | insertion_label = os.path.join(insertion_dir, 'labels') 30 | for img_file in os.listdir(insertion_img): 31 | fname = img_file.split('.png')[0] 32 | img_path = os.path.join(insertion_img, img_file) 33 | shutil.copyfile(img_path, os.path.join(out_img_dir, fname + '_insert.png')) 34 | for txt_file in os.listdir(insertion_label): 35 | fname = txt_file.split('.txt')[0] 36 | txt_path = os.path.join(insertion_label, txt_file) 37 | shutil.copyfile(txt_path, os.path.join(out_label_dir, fname + '_insert.txt')) 38 | 39 | remove_img = os.path.join(remove_dir, 'images') 40 | remove_label = os.path.join(remove_dir, 'labels') 41 | for img_file in os.listdir(remove_img): 42 | fname = img_file.split('.png')[0] 43 | img_path = os.path.join(remove_img, img_file) 44 | shutil.copyfile(img_path, os.path.join(out_img_dir, fname + '_remove.png')) 45 | for txt_file in os.listdir(remove_label): 46 | fname = txt_file.split('.txt')[0] 47 | txt_path = os.path.join(remove_label, txt_file) 48 | shutil.copyfile(txt_path, os.path.join(out_label_dir, fname + '_remove.txt')) 49 | 50 | replace_img = os.path.join(replace_dir, 'images') 51 | replace_label = os.path.join(replace_dir, 'labels') 52 | for img_file in os.listdir(replace_img): 53 | fname = img_file.split('.png')[0] 54 | img_path = os.path.join(replace_img, img_file) 55 | shutil.copyfile(img_path, os.path.join(out_img_dir, fname + '_replace.png')) 56 | for txt_file in os.listdir(replace_label): 57 | fname = txt_file.split('.txt')[0] 58 | txt_path = os.path.join(replace_label, txt_file) 59 | shutil.copyfile(txt_path, os.path.join(out_label_dir, fname + '_replace.txt')) 60 | 61 | 62 | def get_select(path_file, label_in_path, img_out_path, label_out_path, m=100): 63 | file_li = np.load(path_file) 64 | print(file_li[0]) 65 | source_paths, txt_paths, fnames = [], [], [] 66 | for i in range(len(file_li)): 67 | source_path = file_li[i] 68 | fname = source_path.split('\\')[-1].split('.')[0] 69 | if os.path.exists(os.path.join(label_in_path, fname + '.txt')): # 找到匹配的 jpg 和 txt 70 | source_paths.append(source_path) 71 | txt_paths.append(os.path.join(label_in_path, fname + '.txt')) 72 | fnames.append(fname) 73 | print(len(source_paths), len(fnames)) 74 | for i in range(m): 75 | shutil.copyfile(source_paths[i], os.path.join(img_out_path, fnames[i] + '.jpg')) 76 | shutil.copyfile(txt_paths[i], os.path.join(label_out_path, fnames[i] + '.txt')) 77 | 78 | 79 | if __name__ == '__main__': 80 | # # 从 Tsample 中随机选择 100个用例 81 | # img_sample_path = 'E:/code/datasets/coco1000/images/val2017' 82 | # label_sample_path = 'E:/code/datasets/coco1000/labels/val2017' 83 | # img_sample100_path = 'E:/code/datasets/obj_aug_coco1000/select/sample/images' 84 | # label_sample100_path = 'E:/code/datasets/obj_aug_coco1000/select/sample/labels' 85 | # os.makedirs(img_sample100_path, exist_ok=True) 86 | # os.makedirs(label_sample100_path, exist_ok=True) 87 | # random_sample(img_sample_path, label_sample_path, img_sample100_path, label_sample100_path) 88 | 89 | # # 合并三种扩增方式到 Taug 90 | # insertion_dir = 'E:/code/datasets/coco_test_1000/aug/insertion' 91 | # remove_dir = 'E:/code/datasets/coco_test_1000/aug/remove' 92 | # replace_dir = 'E:/code/datasets/coco_test_1000/aug/replace' 93 | # out_dir = 'E:/code/datasets/coco_test_1000/aug/all' 94 | # os.makedirs(out_dir, exist_ok=True) 95 | # get_aug(insertion_dir, remove_dir, replace_dir, out_dir) 96 | 97 | # # 从 Taug 中随机选择 100 个测试用例 98 | # img_aug_path = 'E:/code/datasets/coco1000/aug/all/images' 99 | # label_aug_path = 'E:/code/datasets/coco1000/aug/all/labels' 100 | # img_random_path = 'E:/code/datasets/coco1000/select/random/images' 101 | # label_random_path = 'E:/code/datasets/coco1000/select/random/labels' 102 | # os.makedirs(img_random_path, exist_ok=True) 103 | # os.makedirs(label_random_path, exist_ok=True) 104 | # random_sample(img_aug_path, label_aug_path, img_random_path, label_random_path, m=1000) 105 | 106 | # 获取按照 mAP从小到达选择的测试用例 107 | file_path = 'E:/code/yolov5-master/save_paths.npy' 108 | label_path = 'E:/code/datasets/coco1000/aug/all/labels' 109 | img_select_path = 'E:/code/datasets/coco1000/select/select_20/images' 110 | label_select_path = 'E:/code/datasets/coco1000/select/select_20/labels' 111 | os.makedirs(img_select_path, exist_ok=True) 112 | os.makedirs(label_select_path, exist_ok=True) 113 | get_select(file_path, label_path, img_select_path, label_select_path, m=600) 114 | -------------------------------------------------------------------------------- /utils/yolov5_2_coco.py: -------------------------------------------------------------------------------- 1 | # -*- encoding: utf-8 -*- 2 | # @File: yolov5_2_coco.py 3 | # @Author: SWHL 4 | # @Contact: liekkaskono@163.com 5 | import argparse 6 | import json 7 | import os 8 | import shutil 9 | from pathlib import Path 10 | import time 11 | import warnings 12 | 13 | import cv2 14 | from tqdm import tqdm 15 | 16 | 17 | class YOLOV5ToCOCO(object): 18 | def __init__(self, root_path, txt_path, label_path, class_path): 19 | self.raw_data_dir = Path(root_path) 20 | self.txt_path = Path(txt_path) 21 | self.label_path = Path(label_path) 22 | self.class_path = Path(class_path) 23 | self.mode = 'val' 24 | 25 | # self.verify_exists(self.raw_data_dir / 'images') 26 | # self.verify_exists(self.raw_data_dir / 'labels') 27 | 28 | save_dir_name = f'{Path(self.raw_data_dir).name}_COCO_format' 29 | self.output_dir = self.raw_data_dir.parent / save_dir_name 30 | self.mkdir(self.output_dir) 31 | 32 | self._init_json() 33 | 34 | def __call__(self): 35 | # Read the image txt. 36 | # txt_path = self.raw_data_dir / f'{mode}2017.txt' 37 | self.verify_exists(self.txt_path) 38 | img_list = self.read_txt(self.txt_path) 39 | # if mode == 'train': 40 | # img_list = self.append_bg_img(img_list) 41 | 42 | # Create the directory of saving the new image. 43 | save_img_dir = self.output_dir / f'{self.mode}2017' 44 | self.mkdir(save_img_dir) 45 | 46 | # Generate json file. 47 | anno_dir = self.output_dir / "annotations" 48 | self.mkdir(anno_dir) 49 | 50 | save_json_path = anno_dir / f'instances_{self.mode}2017.json' 51 | json_data = self.convert(img_list, save_img_dir, self.mode) 52 | 53 | self.write_json(save_json_path, json_data) 54 | 55 | def _init_json(self): 56 | # classes_path = self.raw_data_dir / 'classes.txt' 57 | self.verify_exists(self.class_path) 58 | self.categories = self._get_category(self.class_path) 59 | 60 | self.type = 'instances' 61 | self.annotation_id = 1 62 | 63 | self.cur_year = time.strftime('%Y', time.localtime(time.time())) 64 | self.info = { 65 | 'year': int(self.cur_year), 66 | 'version': '1.0', 67 | 'description': 'For object detection', 68 | 'date_created': self.cur_year, 69 | } 70 | 71 | self.licenses = [{ 72 | 'id': 1, 73 | 'name': 'Apache License v2.0', 74 | 'url': 'https://github.com/RapidAI/YOLO2COCO/LICENSE', 75 | }] 76 | 77 | def append_bg_img(self, img_list): 78 | bg_dir = self.raw_data_dir / 'background_images' 79 | if bg_dir.exists(): 80 | bg_img_list = list(bg_dir.iterdir()) 81 | for bg_img_path in bg_img_list: 82 | img_list.append(str(bg_img_path)) 83 | return img_list 84 | 85 | def _get_category(self, classes_path): 86 | class_list = self.read_txt(classes_path) 87 | categories = [] 88 | for i, category in enumerate(class_list, 1): 89 | categories.append({ 90 | 'supercategory': category, 91 | 'id': i, 92 | 'name': category, 93 | }) 94 | return categories 95 | 96 | def convert(self, img_list, save_img_dir, mode): 97 | images, annotations = [], [] 98 | for img_id, img_path in enumerate(tqdm(img_list, desc=mode), 1): 99 | # TODO 100 | print(img_path) 101 | img_path = os.path.join(self.raw_data_dir, img_path.split('./')[1]) 102 | image_dict = self.get_image_info(img_path, img_id, save_img_dir) 103 | images.append(image_dict) 104 | 105 | # TODO 106 | label_path = self.label_path / f'{Path(img_path).stem}.txt' 107 | # label_path = self.raw_data_dir / 'labels' / f'{Path(img_path).stem}.txt' 108 | annotation = self.get_annotation(label_path, 109 | img_id, 110 | image_dict['height'], 111 | image_dict['width']) 112 | annotations.extend(annotation) 113 | 114 | json_data = { 115 | 'info': self.info, 116 | 'images': images, 117 | 'licenses': self.licenses, 118 | 'type': self.type, 119 | 'annotations': annotations, 120 | 'categories': self.categories, 121 | } 122 | return json_data 123 | 124 | def get_image_info(self, img_path, img_id, save_img_dir): 125 | img_path = Path(img_path) 126 | self.verify_exists(img_path) 127 | 128 | new_img_name = f'{img_id:012d}.jpg' 129 | save_img_path = save_img_dir / new_img_name 130 | img_src = cv2.imread(str(img_path)) 131 | if img_path.suffix.lower() == ".jpg": 132 | shutil.copyfile(img_path, save_img_path) 133 | else: 134 | cv2.imwrite(str(save_img_path), img_src) 135 | 136 | height, width = img_src.shape[:2] 137 | image_info = { 138 | 'date_captured': self.cur_year, 139 | 'file_name': new_img_name, 140 | 'id': img_id, 141 | 'height': height, 142 | 'width': width, 143 | } 144 | return image_info 145 | 146 | def get_annotation(self, label_path: Path, img_id, height, width): 147 | def get_box_info(vertex_info, height, width): 148 | cx, cy, w, h = [float(i) for i in vertex_info] 149 | 150 | cx = cx * width 151 | cy = cy * height 152 | box_w = w * width 153 | box_h = h * height 154 | 155 | # left top 156 | x0 = max(cx - box_w / 2, 0) 157 | y0 = max(cy - box_h / 2, 0) 158 | 159 | # right bottom 160 | x1 = min(x0 + box_w, width) 161 | y1 = min(y0 + box_h, height) 162 | 163 | segmentation = [[x0, y0, x1, y0, x1, y1, x0, y1]] 164 | bbox = [x0, y0, box_w, box_h] 165 | area = box_w * box_h 166 | return segmentation, bbox, area 167 | 168 | if not label_path.exists(): 169 | annotation = [{ 170 | 'segmentation': [], 171 | 'area': 0, 172 | 'iscrowd': 0, 173 | 'image_id': img_id, 174 | 'bbox': [], 175 | 'category_id': -1, 176 | 'id': self.annotation_id, 177 | }] 178 | self.annotation_id += 1 179 | return annotation 180 | 181 | annotation = [] 182 | label_list = self.read_txt(str(label_path)) 183 | for i, one_line in enumerate(label_list): 184 | label_info = one_line.split(' ') 185 | if len(label_info) < 5: 186 | warnings.warn( 187 | f'The {i+1} line of the {label_path} has been corrupted.') 188 | continue 189 | 190 | category_id, vertex_info = label_info[0], label_info[1:] 191 | segmentation, bbox, area = get_box_info(vertex_info, height, width) 192 | annotation.append({ 193 | 'segmentation': segmentation, 194 | 'area': area, 195 | 'iscrowd': 0, 196 | 'image_id': img_id, 197 | 'bbox': bbox, 198 | 'category_id': int(category_id)+1, 199 | 'id': self.annotation_id, 200 | }) 201 | self.annotation_id += 1 202 | return annotation 203 | 204 | @staticmethod 205 | def read_txt(txt_path): 206 | with open(str(txt_path), 'r', encoding='utf-8') as f: 207 | data = list(map(lambda x: x.rstrip('\n'), f)) 208 | return data 209 | 210 | @staticmethod 211 | def mkdir(dir_path): 212 | Path(dir_path).mkdir(parents=True, exist_ok=True) 213 | 214 | @staticmethod 215 | def verify_exists(file_path): 216 | file_path = Path(file_path) 217 | if not file_path.exists(): 218 | raise FileNotFoundError(f'The {file_path} is not exists!!!') 219 | 220 | @staticmethod 221 | def write_json(json_path, content: dict): 222 | with open(json_path, 'w', encoding='utf-8') as f: 223 | json.dump(content, f, ensure_ascii=False) 224 | 225 | 226 | if __name__ == "__main__": 227 | parser = argparse.ArgumentParser() 228 | parser.add_argument('--root_path', default='E:/code/datasets/coco', type=str, required=True) 229 | parser.add_argument('--txt_path', default='E:/code/datasets/coco/val2017.txt', type=str, required=True) 230 | parser.add_argument('--label_path', default='E:/code/datasets/coco/labels/val2017', type=str, required=True) 231 | parser.add_argument('--class_path', default='E:/code/datasets/coco/classes.txt', type=str, required=True) 232 | args = parser.parse_args() 233 | 234 | # root_path = 'E:/code/datasets/coco' 235 | # txt_path = 'E:/code/datasets/coco/val2017.txt' 236 | # label_path = 'E:/code/datasets/coco/labels/val2017' 237 | 238 | converter = YOLOV5ToCOCO(args.root_path, args.txt_path, args.label_path, args.class_path) 239 | converter() 240 | -------------------------------------------------------------------------------- /baseline2_aug.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import random 4 | import os 5 | import argparse 6 | 7 | 8 | def get_coord(x, y, w, h, w1, h1): 9 | # 边界框反归一化 10 | x_t = x * w1 11 | y_t = y * h1 12 | w_t = w * w1 13 | h_t = h * h1 14 | 15 | # 计算坐标 16 | top_left_x = x_t - w_t / 2 17 | top_left_y = y_t - h_t / 2 18 | bottom_right_x = x_t + w_t / 2 19 | bottom_right_y = y_t + h_t / 2 20 | 21 | return top_left_x, top_left_y, bottom_right_x, bottom_right_y 22 | 23 | 24 | def get_yolo_label(x1, y1, x2, y2, H, W, mask_class): 25 | new_x = (x1 + x2) / (2 * W) 26 | new_y = (y1 + y2) / (2 * H) 27 | 28 | new_w = (x2 - x1) / W 29 | new_h = (y2 - y1) / H 30 | 31 | # yolo_label_info = str(mask_class) + ' ' + new_x + ' ' + new_y + ' ' + new_w + ' ' + new_h + '\n' 32 | yolo_label_info = '{} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(str(mask_class), new_x, new_y, new_w, new_h) 33 | return yolo_label_info 34 | 35 | 36 | def get_crop_label(img, max_width, max_hight, label_path): 37 | # 读取 label 38 | with open(label_path, 'r') as f: 39 | lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 40 | 41 | h1, w1 = img.shape[:2] 42 | save_label_li = [] 43 | 44 | for i in range(len(lb)): 45 | label_x = lb[i] 46 | label, x, y, w, h = label_x 47 | 48 | top_left_x, top_left_y, bottom_right_x, bottom_right_y = get_coord(x, y, w, h, w1, h1) 49 | if bottom_right_x < max_width and bottom_right_y < max_hight: #原始边框完全在新的图片范围内 50 | tmp = get_yolo_label(top_left_x, top_left_y, bottom_right_x, bottom_right_y, max_hight, max_width, int(label)) 51 | save_label_li.append(tmp) 52 | 53 | return save_label_li 54 | 55 | 56 | # 随机图片裁剪，删除算子的暴力版本 57 | def crop(img_dir, label_dir, save_img_dir, save_label_dir): 58 | for img_file in os.listdir(img_dir): 59 | fname = img_file.split('.jpg')[0] 60 | if os.path.exists(os.path.join(label_dir, fname + '.txt')): # 找到匹配的 jpg 和 txt 61 | jpg_path = os.path.join(img_dir, img_file) 62 | txt_path = os.path.join(label_dir, fname + '.txt') 63 | img = cv2.imread(str(jpg_path)) 64 | 65 | hight, width = img.shape[0], img.shape[1] 66 | random_width = random.randint(1, width) 67 | random_hight = random.randint(1, hight) 68 | cropped = img[0:random_hight, 0:random_width] 69 | img_save_path = os.path.join(save_img_dir, fname + '.png') 70 | cv2.imwrite(img_save_path, cropped) # 新图片存入图片文件夹 71 | 72 | save_label_li = get_crop_label(img, random_width, random_hight, txt_path) 73 | with open(os.path.join(save_label_dir, fname + '.txt'), "w") as f: # 保存新的 label txt 74 | f.writelines(save_label_li) 75 | 76 | 77 | def integration_label(img1, img2, label_path1, label_path2, H, W): # 1在左，2在右 78 | new_label_li = [] 79 | h1, w1 = img1.shape[:2] 80 | with open(label_path1, 'r') as f: 81 | lb1 = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 82 | for i in range(len(lb1)): 83 | label_x = lb1[i] 84 | label, x, y, w, h = label_x 85 | top_left_x1, top_left_y1, bottom_right_x1, bottom_right_y1 = get_coord(x, y, w, h, w1, h1) 86 | tmp = get_yolo_label(top_left_x1, top_left_y1, bottom_right_x1, bottom_right_y1, H, W, int(label)) 87 | new_label_li.append(tmp) 88 | 89 | h2, w2 = img2.shape[:2] 90 | with open(label_path2, 'r') as f: 91 | lb2 = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 92 | for j in range(len(lb2)): 93 | label_x = lb2[j] 94 | label, x, y, w, h = label_x 95 | top_left_x2, top_left_y2, bottom_right_x2, bottom_right_y2 = get_coord(x, y, w, h, w2, h2) 96 | tmp = get_yolo_label(top_left_x2 + w1, top_left_y2, bottom_right_x2 + w1, bottom_right_y2, H, W, int(label)) 97 | new_label_li.append(tmp) 98 | return new_label_li 99 | 100 | 101 | # 随机图片拼贴， 增加算子的暴力版本 102 | def integration(img_dir, label_dir, save_img_dir, save_label_dir): 103 | img_dir_li = os.listdir(img_dir) 104 | for i in range(len(img_dir_li)-1): 105 | img_file1, img_file2 = img_dir_li[i], img_dir_li[i+1] 106 | source_path1, source_path2 = os.path.join(img_dir, img_file1), os.path.join(img_dir, img_file2) 107 | img1, img2 = cv2.imread(source_path1), cv2.imread(source_path2) 108 | fname1, fname2 = img_file1.split('.jpg')[0], img_file2.split('.jpg')[0] 109 | 110 | if os.path.exists(os.path.join(label_dir, fname1 + '.txt')) and os.path.exists(os.path.join(label_dir, fname2 + '.txt')): # 找到匹配的 jpg 和 txt 111 | txt_path1 = os.path.join(label_dir, fname1 + '.txt') 112 | txt_path2 = os.path.join(label_dir, fname2 + '.txt') 113 | 114 | img1_h, img2_h = img1.shape[0], img2.shape[0] 115 | img1_w, img2_w = img1.shape[1], img2.shape[1] 116 | W = img1_w + img2_w 117 | 118 | if img1_h > img2_h: 119 | img2_resize = cv2.copyMakeBorder(img2, 0, img1_h - img2_h, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0)) 120 | res = cv2.hconcat([img1, img2_resize]) # 水平拼接 121 | new_label_li = integration_label(img1, img2, txt_path1, txt_path2, img1_h, W) 122 | elif img2_h > img1_h: 123 | img1_resize = cv2.copyMakeBorder(img1, 0, img2_h - img1_h, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0)) 124 | res = cv2.hconcat([img2, img1_resize]) # 水平拼接 125 | new_label_li = integration_label(img2, img1, txt_path2, txt_path1, img2_h, W) 126 | else: 127 | res = cv2.hconcat([img1, img2]) # 水平拼接 128 | new_label_li = integration_label(img1, img2, txt_path1, txt_path2, img1_h, W) 129 | 130 | img_save_path = os.path.join(save_img_dir, fname1 + '.png') 131 | cv2.imwrite(img_save_path, res) # 新图片存入图片文件夹 132 | 133 | with open(os.path.join(save_label_dir, fname1 + '.txt'), "w") as f: # 保存新的 label txt 134 | f.writelines(new_label_li) 135 | 136 | 137 | def add_mask(image_path, label_path, img_save_path, mask_id): 138 | # 读取图像文件 139 | img = cv2.imread(str(image_path)) 140 | # 读取 label 141 | with open(label_path, 'r') as f: 142 | lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 143 | label_x = lb[mask_id] # 第 mask_id个 object label对应的位置信息 144 | label, x, y, w, h = label_x 145 | h1, w1 = img.shape[:2] 146 | top_left_x, top_left_y, bottom_right_x, bottom_right_y = get_coord(x, y, w, h, w1, h1) 147 | black_color = (0, 0, 0) 148 | cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), 149 | color=black_color, thickness=-1) 150 | cv2.imwrite(img_save_path, img) 151 | 152 | 153 | # 选中一个目标物体直接贴黑块，替换算子的暴力版本 154 | def mask(img_dir, label_dir, save_img_dir, save_label_dir): 155 | for img_file in os.listdir(img_dir): 156 | fname = img_file.split('.jpg')[0] 157 | if os.path.exists(os.path.join(label_dir, fname + '.txt')): # 找到匹配的 jpg 和 txt 158 | jpg_path = os.path.join(img_dir, img_file) 159 | txt_path = os.path.join(label_dir, fname + '.txt') 160 | with open(txt_path, 'r') as f: 161 | lines = f.readlines() 162 | 163 | img_save_path = os.path.join(save_img_dir, fname + '.png') 164 | if len(lines) > 1: # label 框多于1个的时候才做删除，不然就变成没有label的图片了 165 | add_mask(jpg_path, txt_path, img_save_path, mask_id=0) 166 | with open(os.path.join(save_label_dir, fname + '_mask.txt'), "w") as f: # 保存新的 label txt 167 | f.writelines(lines[1:]) 168 | else: # 查找失败，没有可以替换的图片，直接把原始图片和label保存到对应的文件夹 169 | img_save_path = os.path.join(save_img_dir, fname + '.png') 170 | img = cv2.imread(str(jpg_path)) 171 | cv2.imwrite(img_save_path, img) # 保存原始的图片 172 | with open(os.path.join(save_label_dir, fname + '.txt'), "w") as f: # 保存 label txt 173 | f.writelines(lines) 174 | 175 | 176 | if __name__ == '__main__': 177 | parser = argparse.ArgumentParser() 178 | parser.add_argument('--img_dir', type=str, default='../datasets/coco1000/images/val2017', help='img path') 179 | parser.add_argument('--label_dir', type=str, default='../datasets/coco1000/labels/val2017', help='label path') 180 | parser.add_argument('--save_dir', type=str, default='../datasets/obj_aug_coco1000', help='save path') 181 | arg = parser.parse_args() 182 | 183 | img_dir, label_dir, save_dir = arg.img_dir, arg.label_dir, arg.save_dir 184 | crop_save_img_dir = os.path.join(save_dir, 'crop/images') 185 | crop_save_label_dir = os.path.join(save_dir, 'crop/labels') 186 | integration_save_img_dir = os.path.join(save_dir, 'integration/images') 187 | integration_save_label_dir = os.path.join(save_dir, 'integration/labels') 188 | mask_save_img_dir = os.path.join(save_dir, 'mask/images') 189 | mask_save_label_dir = os.path.join(save_dir, 'mask/labels') 190 | 191 | os.makedirs(crop_save_img_dir, exist_ok=True) 192 | os.makedirs(crop_save_label_dir, exist_ok=True) 193 | os.makedirs(integration_save_img_dir, exist_ok=True) 194 | os.makedirs(integration_save_label_dir, exist_ok=True) 195 | os.makedirs(mask_save_img_dir, exist_ok=True) 196 | os.makedirs(mask_save_label_dir, exist_ok=True) 197 | 198 | crop(img_dir, label_dir, crop_save_img_dir, crop_save_label_dir) 199 | integration(img_dir, label_dir, integration_save_img_dir, integration_save_label_dir) 200 | mask(img_dir, label_dir, mask_save_img_dir, mask_save_label_dir) 201 | -------------------------------------------------------------------------------- /img_add.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import math 4 | import random 5 | import os 6 | import argparse 7 | from class_info import coco_class_dict, openimages_class_dict, voc_class_dict 8 | 9 | 10 | def add_alpha_channel(img): 11 | """ 为jpg图像添加alpha通道 """ 12 | 13 | b_channel, g_channel, r_channel = cv2.split(img) # 剥离jpg图像通道 14 | alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255 # 创建Alpha通道 15 | 16 | img_new = cv2.merge((b_channel, g_channel, r_channel, alpha_channel)) # 融合通道 17 | return img_new 18 | 19 | 20 | def add_mask(source_path, mask_path, mask_class, resize_rate=0.05, random_pos=True, pos_h=None, pos_w=None): 21 | mask_img = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) 22 | source_img = cv2.imread(source_path, cv2.IMREAD_UNCHANGED) 23 | 24 | h, w = mask_img.shape[0], mask_img.shape[1] 25 | H, W = source_img.shape[0], source_img.shape[1] 26 | k = w/h 27 | 28 | if source_img.ndim == 2: # 灰度图片先转成三通道的图片 29 | source_img = cv2.cvtColor(source_img, cv2.COLOR_GRAY2RGB) 30 | 31 | # 判断jpg图像是否已经为4通道 32 | if source_img.shape[2] == 3: 33 | source_img = add_alpha_channel(source_img) 34 | 35 | # 面积比 (w * (w/k))/ (W*H) <= resize_rate 36 | resize_w = int(math.sqrt(resize_rate * H * W * k)) if int(math.sqrt(resize_rate * H * W * k)) >= 1 else 1 37 | resize_h = int(resize_w/k) if int(resize_w/k) >= 1 else 1 # 不能变成 0了，至少要是1 38 | resized_mask_img = cv2.resize(mask_img, dsize=(resize_w, resize_h)) 39 | if random_pos: # pos是mask的左上角坐标，只要长宽不超出原图的边界即可 40 | pos_h = random.randint(0, H-resize_h) 41 | pos_w = random.randint(0, W-resize_w) 42 | else: # 固定 mask 的位置，需要检查是否超出边界 43 | assert pos_w is not None 44 | assert pos_h is not None 45 | resize_w = W - pos_w if (pos_w + resize_w > W) else resize_w 46 | resize_h = H - pos_h if (pos_h + resize_h > H) else resize_h 47 | 48 | # 获取要覆盖图像的alpha值，将像素值除以255，使值保持在0-1之间 49 | alpha_png = resized_mask_img[0:resize_h, 0:resize_w, 3] / 255.0 50 | alpha_jpg = 1 - alpha_png 51 | # 开始叠加 52 | for c in range(0, 3): 53 | source_img[pos_h:(pos_h + resize_h), pos_w:(pos_w + resize_w), c] = ( 54 | (alpha_jpg * source_img[pos_h:(pos_h + resize_h), pos_w:(pos_w + resize_w), c]) + (alpha_png * resized_mask_img[0:resize_h, 0:resize_w, c])) 55 | # print("resize mask w:", resize_w, "resize mask h:", resize_h, "\npos w:", pos_w, "pos h:", pos_h) 56 | # print("source w:", W, "source h:", H) 57 | yolo_label_info = yolo_label(pos_w, pos_h, pos_w + resize_w, pos_h + resize_h, H, W, mask_class) 58 | return source_img, yolo_label_info 59 | 60 | 61 | def yolo_label(x1, y1, x2, y2, H, W, mask_class): 62 | new_x = (x1 + x2) / (2 * W) 63 | new_y = (y1 + y2) / (2 * H) 64 | 65 | new_w = (x2 - x1) / W 66 | new_h = (y2 - y1) / H 67 | 68 | # yolo_label_info = str(mask_class) + ' ' + new_x + ' ' + new_y + ' ' + new_w + ' ' + new_h + '\n' 69 | yolo_label_info = '{} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(str(mask_class), new_x, new_y, new_w, new_h) 70 | return yolo_label_info 71 | 72 | 73 | def get_mask_pos(image_path, x): 74 | img = cv2.imread(str(image_path)) 75 | h1, w1 = img.shape[:2] 76 | label, x, y, w, h = x 77 | print("原图宽高:\nw1={}\nh1={}".format(w1, h1)) 78 | 79 | # 边界框反归一化 80 | x_t = x * w1 81 | y_t = y * h1 82 | w_t = w * w1 83 | h_t = h * h1 84 | 85 | # 计算坐标 86 | top_left_x = x_t - w_t / 2 87 | top_left_y = y_t - h_t / 2 88 | bottom_right_x = x_t + w_t / 2 89 | bottom_right_y = y_t + h_t / 2 90 | 91 | mask_area = (bottom_right_x - top_left_x) * (bottom_right_y - top_left_y) 92 | size_ratio = mask_area / (w1 * h1) 93 | print("size_ratio:", size_ratio) 94 | 95 | return int(top_left_x), int(top_left_y), size_ratio 96 | 97 | 98 | def gen_mask(img_dir, label_dir, mask_dir, save_img_dir, save_label_dir): # insertion 99 | for img_file in os.listdir(img_dir): 100 | fname = img_file.split('.jpg')[0] 101 | print(fname) 102 | if os.path.exists(os.path.join(label_dir, fname + '.txt')): # 找到匹配的 jpg 和 txt 103 | source_path = os.path.join(img_dir, img_file) 104 | txt_path = os.path.join(label_dir, fname + '.txt') 105 | with open(txt_path, 'r') as f: 106 | lines = f.readlines() 107 | random_mask = random.sample(os.listdir(mask_dir), 1) 108 | mask_path = os.path.join(mask_dir, random_mask[0]) 109 | mask_class = random_mask[0].split('_')[0] 110 | 111 | added_mask_img, mask_yolo_label = add_mask(source_path, mask_path, mask_class) 112 | lines.append(mask_yolo_label) # 在原始 label的基础上把 mask的 label加上 113 | img_save_path = os.path.join(save_img_dir, fname + '.png') 114 | cv2.imwrite(img_save_path, added_mask_img) # 增加了mask的图片存入图片文件夹 115 | with open(os.path.join(save_label_dir, fname + '.txt'), "w") as f: # 增加了 mask的新 label存入txt文件夹 116 | f.writelines(lines) 117 | 118 | 119 | def replace_mark(img_dir, label_dir, mask_dir, save_img_dir, save_label_dir, dataset): 120 | for img_file in os.listdir(img_dir): 121 | fname = img_file.split('_mask.png')[0] 122 | print(fname) 123 | if os.path.exists(os.path.join(label_dir, fname + '_mask.txt')): # 找到匹配的 jpg 和 txt 124 | source_path = os.path.join(img_dir, img_file) 125 | ori_txt_path = os.path.join(label_dir, fname + '_mask.txt') 126 | replace_txt_path = os.path.join(label_dir, fname + '_replace.txt') 127 | with open(ori_txt_path, 'r') as f1: 128 | lines = f1.readlines() 129 | with open(replace_txt_path, 'r') as f2: 130 | line = f2.readline() 131 | lb = [float(i) for i in line.strip().split(' ')] 132 | replace_class = int(lb[0]) 133 | if dataset == 'coco': 134 | mask_path = os.path.join(mask_dir, coco_class_dict[replace_class]) 135 | elif dataset == 'openimage': 136 | mask_path = os.path.join(mask_dir, openimages_class_dict[replace_class]) 137 | elif dataset == 'voc': 138 | mask_path = os.path.join(mask_dir, voc_class_dict[replace_class]) 139 | pos_w, pos_h, size_ratio = get_mask_pos(source_path, lb) 140 | 141 | added_mask_img, mask_yolo_label = add_mask(source_path, mask_path, replace_class, resize_rate=size_ratio, 142 | random_pos=False, pos_h=pos_h, pos_w=pos_w) 143 | lines.append(mask_yolo_label) # 在原始 label的基础上把 mask的 label加上 144 | img_save_path = os.path.join(save_img_dir, fname + '.png') 145 | cv2.imwrite(img_save_path, added_mask_img) # 增加了mask的图片存入图片文件夹 146 | with open(os.path.join(save_label_dir, fname + '.txt'), "w") as f: # 增加了 mask的新 label存入txt文件夹 147 | f.writelines(lines) 148 | 149 | 150 | def insertion(dataset, type): 151 | if dataset == 'coco' and type == 'val': 152 | mask_dir = "./coco_classes" 153 | img_dir = "../datasets/coco1000/images/val2017" 154 | label_dir = "../datasets/coco1000/labels/val2017" 155 | save_img_dir = "../datasets/obj_aug_coco1000/insertion/images" 156 | save_label_dir = "../datasets/obj_aug_coco1000/insertion/labels" 157 | 158 | elif dataset == 'coco' and type == 'train': 159 | mask_dir = "./coco_classes" 160 | img_dir = "../datasets/coco1000/images/train2017" 161 | label_dir = "../datasets/coco1000/labels/train2017" 162 | save_img_dir = "../datasets/coco1000/images/train2017_aug/insertion" 163 | save_label_dir = "../datasets/coco1000/labels/train2017_aug/insertion" 164 | 165 | elif dataset == 'openimage' and type == 'val': 166 | mask_dir = "./openimages_classes" 167 | img_dir = "E:/code/mmdetection/data/OpenImages1000/OpenImages/validation" 168 | label_dir = "E:/code/mmdetection/data/OpenImages1000/labels" 169 | save_img_dir = "../datasets/obj_aug_openimages1000/insertion/images" 170 | save_label_dir = "../datasets/obj_aug_openimages1000/insertion/labels" 171 | 172 | elif dataset == 'voc' and type == 'val': 173 | mask_dir = "./voc_classes" 174 | img_dir = "E:/code/mmdetection/data/VOCdevkit/test/VOC2007/JPEGImages1000" 175 | label_dir = "E:/code/mmdetection/data/VOCdevkit/test/labels" 176 | save_img_dir = "../datasets/obj_aug_voc1000/insertion/images" 177 | save_label_dir = "../datasets/obj_aug_voc1000/insertion/labels" 178 | 179 | os.makedirs(save_img_dir, exist_ok=True) 180 | os.makedirs(save_label_dir, exist_ok=True) 181 | 182 | gen_mask(img_dir, label_dir, mask_dir, save_img_dir, save_label_dir) 183 | 184 | 185 | def replace(dataset, type): 186 | if dataset == 'coco' and type == 'val': 187 | mask_dir = "./coco_classes" 188 | img_dir = "../datasets/coco1000/images/train2017_aug/replace/images_tmp2" 189 | label_dir = "../datasets/coco1000/labels/train2017_aug/replace_tmp" 190 | save_img_dir = "../datasets/coco1000/images/train2017_aug/replace" 191 | save_label_dir = "../datasets/coco1000/labels/train2017_aug/replace" 192 | elif dataset == 'coco' and type == 'train': 193 | mask_dir = "./coco_classes" 194 | img_dir = "../datasets/obj_aug_coco1000/replace/images_tmp2" 195 | label_dir = "../datasets/obj_aug_coco1000/replace/labels_tmp" 196 | save_img_dir = "../datasets/obj_aug_coco1000/replace/images" 197 | save_label_dir = "../datasets/obj_aug_coco1000/replace/labels" 198 | elif dataset == 'openimage' and type == 'val': 199 | mask_dir = "./openimages_classes" 200 | img_dir = "../datasets/obj_aug_openimages1000/replace/images_tmp2" 201 | label_dir = "../datasets/obj_aug_openimages1000/replace/labels_tmp" 202 | save_img_dir = "../datasets/obj_aug_openimages1000/replace/images" 203 | save_label_dir = "../datasets/obj_aug_openimages1000/replace/labels" 204 | elif dataset == 'voc' and type == 'val': 205 | mask_dir = "./voc_classes" 206 | img_dir = "../datasets/obj_aug_voc1000/replace/images_tmp2" 207 | label_dir = "../datasets/obj_aug_voc1000/replace/labels_tmp" 208 | save_img_dir = "../datasets/obj_aug_voc1000/replace/images" 209 | save_label_dir = "../datasets/obj_aug_voc1000/replace/labels" 210 | 211 | os.makedirs(save_img_dir, exist_ok=True) 212 | os.makedirs(save_label_dir, exist_ok=True) 213 | 214 | replace_mark(img_dir, label_dir, mask_dir, save_img_dir, save_label_dir, dataset) 215 | 216 | 217 | if __name__ == '__main__': 218 | parser = argparse.ArgumentParser() 219 | parser.add_argument('--task', required=True, choices=['insertion', 'replace']) 220 | parser.add_argument('--dataset', required=True, choices=['coco', 'voc', 'openimage']) 221 | parser.add_argument('--type', default='val', required=True, choices=['val', 'train']) 222 | args = parser.parse_args() 223 | 224 | if args.task == 'insertion': 225 | insertion(args.dataset, args.type) 226 | if args.task == 'replace': 227 | replace(args.dataset, args.type) 228 | -------------------------------------------------------------------------------- /gen_mask_img.py: -------------------------------------------------------------------------------- 1 | # 读取coco图片，选定一个框进行mask，然后保存对应的 mask png图片，以及对应新的label 2 | import cv2 3 | import numpy as np 4 | import os 5 | import argparse 6 | 7 | 8 | coco_clean_classes = {0, 2, 4, 6, 8, 10, 11, 15, 16, 22, 23, 46, 48, 49, 53, 54, 61, 72, 74} 9 | openimage_clean_classes = {110, 121, 228, 34, 404, 426, 433, 440, 445, 455, 467, 483, 498, 503, 510, 570, 592, 595, 68, 98} 10 | voc_clean_classes = {0, 1, 2, 4, 5, 6, 7, 11, 14, 18} 11 | 12 | 13 | # (w1,h1)是原始图片的尺寸，x包含 object label的位置信息 14 | def mask(img, x, img_save_path, mask_save_path): 15 | cv2.imwrite(img_save_path, img) 16 | cv2.destroyAllWindows() 17 | 18 | h1, w1 = img.shape[:2] 19 | label, x, y, w, h = x 20 | print("原图宽高:\nw1={}\nh1={}".format(w1, h1)) 21 | 22 | # 边界框反归一化 23 | x_t = x * w1 24 | y_t = y * h1 25 | w_t = w * w1 26 | h_t = h * h1 27 | 28 | # 计算坐标 29 | top_left_x = x_t - w_t / 2 30 | top_left_y = y_t - h_t / 2 31 | bottom_right_x = x_t + w_t / 2 32 | bottom_right_y = y_t + h_t / 2 33 | print("左上x坐标:{}".format(top_left_x)) 34 | print("左上y坐标:{}".format(top_left_y)) 35 | print("右下x坐标:{}".format(bottom_right_x)) 36 | print("右下y坐标:{}".format(bottom_right_y)) 37 | 38 | black = np.zeros((h1, w1)) # 全黑底色图片，与原图的尺寸一样大 39 | white_color = (255, 255, 255) 40 | cv2.rectangle(black, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), 41 | color=white_color, thickness=-1) 42 | # cv2.imshow('show', black) 43 | cv2.imwrite(mask_save_path, black) 44 | # cv2.waitKey(0) # 按键结束 45 | cv2.destroyAllWindows() 46 | 47 | 48 | def gen_mask(image_path, label_path, img_save_path, mask_save_path, random_mask, mask_id): 49 | # 读取图像文件 50 | img = cv2.imread(str(image_path)) 51 | # 读取 label 52 | with open(label_path, 'r') as f: 53 | lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 54 | 55 | if random_mask: 56 | area_li = [] 57 | for i in range(len(lb)): 58 | x = lb[i] 59 | label, x, y, w, h = x 60 | area_li.append(w * h) 61 | mask_id = area_li.index(min(area_li)) 62 | x = lb[mask_id] # 第一个 object label对应的位置信息 63 | mask(img, x, img_save_path, mask_save_path) 64 | else: 65 | x = lb[mask_id] # 第一个 object label对应的位置信息 66 | mask(img, x, img_save_path, mask_save_path) 67 | 68 | 69 | # 存储每张图片的 mask，以及增加 mask之后的 label txt文件 70 | def get_coco(img_dir, label_dir, save_img_tmp_dir, save_img_dir, save_label_dir): 71 | for img_file in os.listdir(img_dir): 72 | fname = img_file.split('.jpg')[0] 73 | if os.path.exists(os.path.join(label_dir, fname + '.txt')): # 找到匹配的 jpg 和 txt 74 | jpg_path = os.path.join(img_dir, img_file) 75 | txt_path = os.path.join(label_dir, fname + '.txt') 76 | f = open(txt_path) 77 | lines = f.readlines() # 原始 label 78 | # 保存 mask png图片 79 | mask_save_path = os.path.join(save_img_tmp_dir, fname + '_mask.png') 80 | img_save_path = os.path.join(save_img_tmp_dir, fname + '.png') 81 | 82 | if len(lines) > 1: # label 框多于1个的时候才做删除，不然就变成没有label的图片了 83 | gen_mask(jpg_path, txt_path, img_save_path, mask_save_path, random_mask=True, mask_id=-1) 84 | with open(os.path.join(save_label_dir, fname + '_mask.txt'), "w") as f: # 保存新的 label txt 85 | f.writelines(lines[1:]) 86 | else: # label 框只有1个的时候，直接把原始图片和label保存到对应的文件夹 87 | img_save_path = os.path.join(save_img_dir, fname + '.png') 88 | img = cv2.imread(str(jpg_path)) 89 | cv2.imwrite(img_save_path, img) # 保存原始的图片 90 | cv2.destroyAllWindows() 91 | with open(os.path.join(save_label_dir, fname + '.txt'), "w") as f: # 保存 label txt 92 | f.writelines(lines) 93 | 94 | 95 | # 存储每张图片的 mask，以及增加 mask之后的 label txt文件 96 | def get_specific_coco(img_dir, label_dir, save_img_tmp_dir, save_img_dir, save_label_tmp_dir, save_label_dir, dataset): 97 | for img_file in os.listdir(img_dir): 98 | fname = img_file.split('.jpg')[0] 99 | if os.path.exists(os.path.join(label_dir, fname + '.txt')): # 找到匹配的 jpg 和 txt 100 | jpg_path = os.path.join(img_dir, img_file) 101 | txt_path = os.path.join(label_dir, fname + '.txt') 102 | with open(txt_path, 'r') as f: 103 | lines = f.readlines() 104 | obj_class = [int(lines[i].split(' ')[0]) for i in range(len(lines))] # 每个label 对应的 class 105 | print(obj_class) 106 | del_ids = [] 107 | if dataset == 'coco': 108 | del_ids = [idx for (idx, tmp) in enumerate(obj_class) if tmp in coco_clean_classes] 109 | elif dataset == 'openimage': 110 | del_ids = [idx for (idx, tmp) in enumerate(obj_class) if tmp in openimage_clean_classes] 111 | elif dataset == 'voc': 112 | del_ids = [idx for (idx, tmp) in enumerate(obj_class) if tmp in voc_clean_classes] 113 | # 保存 mask png图片 114 | mask_save_path = os.path.join(save_img_tmp_dir, fname + '_mask.png') 115 | img_save_path = os.path.join(save_img_tmp_dir, fname + '.png') 116 | 117 | delete = True 118 | if len(del_ids) > 0: 119 | del_id = del_ids[0] 120 | 121 | with open(txt_path, 'r') as f: 122 | lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels 123 | label, x, y, w, h = lb[del_id] 124 | area = w * h 125 | if area <= 0.5: # 删除的面积占比小于0.5才行 126 | gen_mask(jpg_path, txt_path, img_save_path, mask_save_path, random_mask=False, mask_id=del_id) 127 | with open(os.path.join(save_label_tmp_dir, fname + '_mask.txt'), "w") as f: # 保存新的 label txt 128 | f.writelines(lines[0:del_id] + lines[del_id+1:]) 129 | with open(os.path.join(save_label_tmp_dir, fname + '_replace.txt'), "w") as f: # 保存待替换的label txt 130 | f.writelines(lines[del_id]) 131 | else: 132 | delete = False 133 | if delete == False or len(del_ids) == 0: # 查找失败，没有可以替换的图片，直接把原始图片和label保存到对应的文件夹 134 | img_save_path = os.path.join(save_img_dir, fname + '.png') 135 | img = cv2.imread(str(jpg_path)) 136 | cv2.imwrite(img_save_path, img) # 保存原始的图片 137 | cv2.destroyAllWindows() 138 | with open(os.path.join(save_label_dir, fname + '.txt'), "w") as f: # 保存 label txt 139 | f.writelines(lines) 140 | 141 | 142 | def prepare_remove(dataset, type): 143 | if dataset == 'coco' and type == 'val': 144 | img_dir = "../datasets/coco1000/images/val2017" 145 | label_dir = "../datasets/coco1000/labels/val2017" 146 | save_img_tmp_dir = "../datasets/obj_aug_coco1000/remove/images_tmp" 147 | save_img_dir = "../datasets/obj_aug_coco1000/remove/images" 148 | save_label_dir = "../datasets/obj_aug_coco1000/remove/labels" 149 | 150 | elif dataset == 'coco' and type == 'train': 151 | img_dir = "../datasets/coco1000/images/train2017" 152 | label_dir = "../datasets/coco1000/labels/train2017" 153 | save_img_tmp_dir = "../datasets/coco1000/images/train2017_aug/remove/images_tmp" 154 | save_img_dir = "../datasets/coco1000/images/train2017_aug/remove" 155 | save_label_dir = "../datasets/coco1000/labels/train2017_aug/remove" 156 | 157 | elif dataset == 'openimage' and type == 'val': 158 | img_dir = "E:/code/mmdetection/data/OpenImages1000/OpenImages/validation" 159 | label_dir = "E:/code/mmdetection/data/OpenImages1000/labels" 160 | save_img_tmp_dir = "../datasets/obj_aug_openimages1000/remove/images_tmp" 161 | save_img_dir = "../datasets/obj_aug_openimages1000/remove/images" 162 | save_label_dir = "../datasets/obj_aug_openimages1000/remove/labels" 163 | 164 | elif dataset == 'voc' and type == 'val': 165 | img_dir = "E:/code/mmdetection/data/VOCdevkit/test/VOC2007/JPEGImages1000" 166 | label_dir = "E:/code/mmdetection/data/VOCdevkit/test/labels" 167 | save_img_tmp_dir = "../datasets/obj_aug_voc1000/remove/images_tmp" 168 | save_img_dir = "../datasets/obj_aug_voc1000/remove/images" 169 | save_label_dir = "../datasets/obj_aug_voc1000/remove/labels" 170 | 171 | os.makedirs(save_img_tmp_dir, exist_ok=True) 172 | os.makedirs(save_img_dir, exist_ok=True) 173 | os.makedirs(save_label_dir, exist_ok=True) 174 | 175 | get_coco(img_dir, label_dir, save_img_tmp_dir, save_img_dir, save_label_dir) # 接下来，save_img_dir可以直接做为 lama 的输入文件夹，lama的输出文件夹与save_label_dir组合起来，就可以用于yolo检测验证了 176 | 177 | 178 | def prepare_replace(dataset, type): 179 | if dataset == 'coco' and type == 'val': 180 | img_dir = "../datasets/coco1000/images/val2017" 181 | label_dir = "../datasets/coco1000/labels/val2017" 182 | save_img_tmp_dir = "../datasets/obj_aug_coco1000/replace/images_tmp1" 183 | save_img_dir = "../datasets/obj_aug_coco1000/replace/images" 184 | save_label_tmp_dir = "../datasets/obj_aug_coco1000/replace/labels_tmp" 185 | save_label_dir = "../datasets/obj_aug_coco1000/replace/labels" 186 | 187 | elif dataset == 'coco' and type == 'train': 188 | img_dir = "../datasets/coco1000/images/train2017" 189 | label_dir = "../datasets/coco1000/labels/train2017" 190 | save_img_tmp_dir = "../datasets/coco1000/images/train2017_aug/replace/images_tmp1" 191 | save_img_dir = "../datasets/coco1000/images/train2017_aug/replace" 192 | save_label_tmp_dir = "../datasets/coco1000/labels/train2017_aug/replace_tmp" 193 | save_label_dir = "../datasets/coco1000/labels/train2017_aug/replace" 194 | 195 | elif dataset == 'openimage' and type == 'val': 196 | img_dir = "E:/code/mmdetection/data/OpenImages1000/OpenImages/validation" 197 | label_dir = "E:/code/mmdetection/data/OpenImages1000/labels" 198 | save_img_tmp_dir = "../datasets/obj_aug_openimages1000/replace/images_tmp1" 199 | save_img_dir = "../datasets/obj_aug_openimages1000/replace/images" 200 | save_label_tmp_dir = "../datasets/obj_aug_openimages1000/replace/labels_tmp" 201 | save_label_dir = "../datasets/obj_aug_openimages1000/replace/labels" 202 | 203 | elif dataset == 'voc' and type == 'val': 204 | img_dir = "E:/code/mmdetection/data/VOCdevkit/test/VOC2007/JPEGImages1000" 205 | label_dir = "E:/code/mmdetection/data/VOCdevkit/test/labels" 206 | save_img_tmp_dir = "../datasets/obj_aug_voc1000/replace/images_tmp1" 207 | save_img_dir = "../datasets/obj_aug_voc1000/replace/images" 208 | save_label_tmp_dir = "../datasets/obj_aug_voc1000/replace/labels_tmp" 209 | save_label_dir = "../datasets/obj_aug_voc1000/replace/labels" 210 | 211 | os.makedirs(save_img_tmp_dir, exist_ok=True) 212 | os.makedirs(save_img_dir, exist_ok=True) 213 | os.makedirs(save_label_tmp_dir, exist_ok=True) 214 | os.makedirs(save_label_dir, exist_ok=True) 215 | 216 | get_specific_coco(img_dir, label_dir, save_img_tmp_dir, save_img_dir, save_label_tmp_dir, save_label_dir, dataset) 217 | 218 | 219 | if __name__ == '__main__': 220 | parser = argparse.ArgumentParser() 221 | parser.add_argument('--task', default='val', required=True, choices=['remove', 'replace']) 222 | parser.add_argument('--dataset', required=True, choices=['coco', 'voc', 'openimage']) 223 | parser.add_argument('--type', default='val', required=True, choices=['val', 'train']) 224 | args = parser.parse_args() 225 | 226 | if args.task == 'remove': 227 | prepare_remove(args.dataset, args.type) 228 | if args.task == 'replace': 229 | prepare_replace(args.dataset, args.type) 230 | --------------------------------------------------------------------------------