目标检测数据集小目标扩增，小目标增强

import os import cv2 from tqdm import tqdm import random import numpy as np def check_dir(path): if not os.path.exists(path): os.makedirs(path) def replace_labels(path): label_dir = [] for p in path: p = p.replace('images','labels') label_dir.append(p.replace('.jpg','.txt')) return label_dir def copyso2(base_dir,image_dir, label_dir, save_base_dir, small_img,small_img_path,small_img_labels,c,rs): #原图 image = cv2.imread(image_dir) labels = read_label_txt(label_dir) if len(labels) == 0: return rescale_labels = rescale_yolo_labels(labels, image.shape) #存储之前框 all_boxes = [] for _,rescale_label in enumerate(rescale_labels): all_boxes.append(rescale_label) for small_img_dir in small_img: image_bbox = cv2.imread(small_img_dir) temp ='' for x in small_img_path: if x in small_img_dir: temp = x break if rs: roi = suo_fang(image_bbox,area_max=3000,area_min=1500) else: if image.shape[0]<=image_bbox.shape[0] or image.shape[1]<=image_bbox.shape[1]: roi = suo_fang(image_bbox,area_max=3000,area_min=1500) else: roi = image_bbox new_bboxes = random_add_patches2(roi.shape, rescale_labels, image.shape, temp, small_img_labels,paste_number=1, iou_thresh=0) count =0 for new_bbox in new_bboxes: count += 1 cl, bbox_left, bbox_top, bbox_right, bbox_bottom = new_bbox[0], new_bbox[1], new_bbox[2], new_bbox[3], \ new_bbox[4] #roi = GaussianBlurImg(roi) # 高斯模糊 height, width, channels = roi.shape center = (int(width / 2),int(height / 2)) #ran_point = (int((bbox_top+bbox_bottom)/2),int((bbox_left+bbox_right)/2)) mask = 255 * np.ones(roi.shape, roi.dtype) try: if count > 1: roi = flip_bbox(roi) image[bbox_top:bbox_bottom, bbox_left:bbox_right] = cv2.seamlessClone(roi, image[bbox_top:bbox_bottom, bbox_left:bbox_right], mask, center, cv2.NORMAL_CLONE) all_boxes.append(new_bbox) rescale_labels.append(new_bbox) except ValueError: continue # dir_name = find_str(image_dir) save_dir=save_base_dir check_dir(save_dir) # save_dir = join(save_base_dir, dir_name) if not os.path.exists(os.path.join(save_dir,'labels',c)): os.makedirs(os.path.join(save_dir, "labels", c)) if not os.path.exists(os.path.join(save_dir, "images", c)): os.makedirs(os.path.join(save_dir, "images", c)) yolo_txt_dir = os.path.join(save_dir,'labels',c, os.path.basename(image_dir.replace('.jpg', '_augment.txt'))) cv2.imwrite(os.path.join(save_dir,'images' ,c,os.path.basename(image_dir).replace('.jpg', '_augment.jpg')), image) convert_all_boxes(image.shape, all_boxes, yolo_txt_dir) def copyso(image_dir, label_dir, save_base_dir,c): image = cv2.imread(image_dir) ih ,iw = image.shape[0],image.shape[1] labels = read_label_txt(label_dir) if len(labels) == 0: return rescale_labels = rescale_yolo_labels(labels, image.shape) #存储之前框 all_boxes = [] for _,rescale_label in enumerate(rescale_labels): all_boxes.append(rescale_label) re_boxes = [] for temp in labels: tm = [] a = list((map(float,temp))) cl = int(a[0]) tm.append(cl) w = iw * a[3] h = ih * a[4] xmin = int(float(a[1]*iw-(w/2))) ymin = int(float(a[2]*ih-(h/2))) xmax = int(float(xmin)+float(w)) ymax = int(float(ymin)+float(h)) tm.append(xmin) tm.append(ymin) tm.append(xmax) tm.append(ymax) re_boxes.append(tm) for i in re_boxes: small_img = image[i[2]:i[4],i[1]:i[3]] roi = suo_fang(small_img,area_max=3000,area_min=1500) new_bboxes = random_add_patches(roi.shape, rescale_labels, image.shape,i[0],paste_number=1, iou_thresh=0) count =0 for new_bbox in new_bboxes: count += 1 cl, bbox_left, bbox_top, bbox_right, bbox_bottom = new_bbox[0], new_bbox[1], new_bbox[2], new_bbox[3], \ new_bbox[4] height, width, channels = roi.shape center = (int(width / 2),int(height / 2)) mask = 255 * np.ones(roi.shape, roi.dtype) try: if count > 1: roi = flip_bbox(roi) image[bbox_top:bbox_bottom, bbox_left:bbox_right] = cv2.seamlessClone(roi, image[bbox_top:bbox_bottom, bbox_left:bbox_right], mask, center, cv2.NORMAL_CLONE) all_boxes.append(new_bbox) rescale_labels.append(new_bbox) except ValueError: continue # dir_name = find_str(image_dir) save_dir=save_base_dir check_dir(save_dir) # save_dir = join(save_base_dir, dir_name) if not os.path.exists(os.path.join(save_dir,'labels',c)): os.makedirs(os.path.join(save_dir, "labels", c)) if not os.path.exists(os.path.join(save_dir, "images", c)): os.makedirs(os.path.join(save_dir, "images", c)) yolo_txt_dir = os.path.join(save_dir,'labels',c, os.path.basename(image_dir.replace('.jpg', '_augment.txt'))) cv2.imwrite(os.path.join(save_dir,'images' ,c,os.path.basename(image_dir).replace('.jpg', '_augment.jpg')), image) convert_all_boxes(image.shape, all_boxes, yolo_txt_dir) def convert_all_boxes(shape, anno_infos, yolo_label_txt_dir): height, width, n = shape label_file = open(yolo_label_txt_dir, 'w') for anno_info in anno_infos: target_id, x1, y1, x2, y2 = anno_info b = (float(x1), float(x2), float(y1), float(y2)) bb = convert((width, height), b) label_file.write(str(target_id) + " " + " ".join([str(a) for a in bb]) + '\n') def convert(size, box): dw = 1. / (size[0]) dh = 1. / (size[1]) x = (box[0] + box[1]) / 2.0 - 1 y = (box[2] + box[3]) / 2.0 - 1 w = box[1] - box[0] h = box[3] - box[2] x = x * dw w = w * dw y = y * dh h = h * dh return (x, y, w, h) def read_label_txt(label_dir): labels = [] with open(label_dir) as fp: for f in fp.readlines(): labels.append(f.strip().split(' ')) return labels def flip_bbox(roi): roi = roi[:, ::-1, :] return roi def make_txt(path,flag): images_path = os.path.join(path,'images') ftest = open(os.path.join(path,'test.txt'), 'a+') ftrain = open(os.path.join(path,'train.txt'), 'a+') fval = open(os.path.join(path,'val.txt'), 'a+') xmlfilepath = os.path.join(path,'images','train') txtsavepath = os.path.join(path,'images','val') testsavapath = os.path.join(path,'images','test') total_xml = os.listdir(xmlfilepath) total_val = os.listdir(txtsavepath) total_test = os.listdir(testsavapath) print('train.txt--begin') for i in tqdm(range(len(total_xml))): ftrain.write(xmlfilepath+"\\"+total_xml[i]) ftrain.write('\n') print('val.txt--begin') for i in tqdm(range(len(total_va