def directCreateDataset(gtFile, outputPath, checkValid=True, map_size=8589934592, data_dir=cfg.data_dir): """ Create LMDB dataset for training and evaluation. ARGS: inputPath : input folder path where starts imagePath outputPath : LMDB output path gtFile : list of image path and label checkValid : if true, check the validity of every image """ os.makedirs(outputPath, exist_ok=True) env = lmdb.open(outputPath, map_size=map_size) # 85899345920/8Gb cache = {} cnt = 1 gtFile = os.path.join(data_dir, gtFile) with open(gtFile, 'r') as data: f_list = data.readlines() nSamples = len(f_list) for line, _ in zip(f_list, tqdm(range(nSamples))): print('第{}张图片:{}'.format(cnt, f_list[cnt - 1])) line_cols = str(line).strip().split(',') img_name, width, height = \ line_cols[0].strip(), int(line_cols[1].strip()), \ int(line_cols[2].strip()) gt = np.zeros((height // cfg.pixel_size, width // cfg.pixel_size, 7)) train_label_dir = os.path.join( data_dir, cfg.train_label_dir_name) # 'labels_%s/' % train_task_id xy_list_array = np.load( os.path.join(train_label_dir, img_name[:-4] + '.npy')) # (N, 4, 2) train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) if not os.path.exists(os.path.join(train_image_dir, img_name)): print('%s does not exist' % os.path.join(train_image_dir, img_name)) continue # ---------------------------------生成标签--------------------------------- with Image.open(os.path.join(train_image_dir, img_name)) as im: draw = ImageDraw.Draw(im) for xy_list in xy_list_array: _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) p_min = np.amin(shrink_xy_list, axis=0) p_max = np.amax(shrink_xy_list, axis=0) # floor of the float ji_min = (p_min / cfg.pixel_size - 0.5).astype(int) - 1 # +1 for ceil of the float and +1 for include the end ji_max = (p_max / cfg.pixel_size - 0.5).astype(int) + 3 imin = np.maximum(0, ji_min[1]) imax = np.minimum(height // cfg.pixel_size, ji_max[1]) jmin = np.maximum(0, ji_min[0]) jmax = np.minimum(width // cfg.pixel_size, ji_max[0]) for i in range(imin, imax): for j in range(jmin, jmax): px = (j + 0.5) * cfg.pixel_size py = (i + 0.5) * cfg.pixel_size if point_inside_of_quad(px, py, shrink_xy_list, p_min, p_max): gt[i, j, 0] = 1 line_width, line_color = 1, 'red' ith = point_inside_of_nth_quad( px, py, xy_list, shrink_1, long_edge) vs = [[[3, 0], [1, 2]], [[0, 1], [2, 3]]] if ith in range(2): gt[i, j, 1] = 1 if ith == 0: line_width, line_color = 2, 'yellow' else: line_width, line_color = 2, 'green' gt[i, j, 2:3] = ith gt[i, j, 3:5] = \ xy_list[vs[long_edge][ith][0]] - [px, py] gt[i, j, 5:] = \ xy_list[vs[long_edge][ith][1]] - [px, py] draw.line([(px - 0.5 * cfg.pixel_size, py - 0.5 * cfg.pixel_size), (px + 0.5 * cfg.pixel_size, py - 0.5 * cfg.pixel_size), (px + 0.5 * cfg.pixel_size, py + 0.5 * cfg.pixel_size), (px - 0.5 * cfg.pixel_size, py + 0.5 * cfg.pixel_size), (px - 0.5 * cfg.pixel_size, py - 0.5 * cfg.pixel_size)], width=line_width, fill=line_color) act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if cfg.draw_act_quad: im.save(os.path.join(act_image_dir, img_name)) # train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) # 'labels_%s/' % train_task_id # np.save(os.path.join(train_label_dir, img_name[:-4] + '_gt.npy'), gt) imagePath = os.path.join(cfg.data_dir, cfg.train_image_dir_name, img_name) label = gt # ---------------------------写入LMDB--------------------------- with open(imagePath, 'rb') as f: imageBin = f.read() if checkValid: try: if not checkImageIsValid(imageBin): print('%s is not a valid image' % imagePath) continue except (Exception): print('error occured', i) with open(outputPath + '/error_image_log.txt', 'a') as log: log.write('%s-th image data occured error\n' % str(i)) continue imageKey = 'image-%09d'.encode() % cnt labelKey = 'label-%09d'.encode() % cnt gt_xy_list_Key = 'gt_xy_list-%09d'.encode() % cnt cache[imageKey] = imageBin cache[labelKey] = label cache[gt_xy_list_Key] = xy_list_array if cnt % 1000 == 0: writeCache(env, cache) cache = {} print('Written %d / %d' % (cnt, nSamples)) cnt += 1 nSamples = cnt - 1 cache['num-samples'.encode()] = str(nSamples).encode() cache['width-height'.encode()] = str(width).encode() writeCache(env, cache) print('Created dataset with %d samples' % nSamples)
def preprocess(): data_dir = cfg.data_dir origin_image_dir = os.path.join(data_dir, cfg.origin_image_dir_name) origin_txt_dir = os.path.join(data_dir, cfg.origin_txt_dir_name) train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) draw_gt_quad = cfg.draw_gt_quad show_gt_image_dir = os.path.join(data_dir, cfg.show_gt_image_dir_name) if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) show_act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) o_img_list = os.listdir(origin_image_dir) print('found %d origin images.' % len(o_img_list)) train_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # d_wight, d_height = resize_image(im) d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') show_gt_im = im.copy() # draw on the img draw = ImageDraw.Draw(show_gt_im) with open( os.path.join(origin_txt_dir, 'object_' + o_img_fname[:-4]), 'r') as f: anno_list = f.readlines() xy_list_array = np.zeros((len(anno_list), 5, 2)) for anno, i in zip(anno_list, range(len(anno_list))): category, anno = anno.strip().split('\t') anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h xy_list = reorder_vertexes(xy_list) # set the class label, needed one-hot encoding if category == 'circle': xy_list_array[i] = np.r_[xy_list, np.array([[0, 1]])] elif category == 'rect': xy_list_array[i] = np.r_[xy_list, np.array([[1, 0]])] else: xy_list_array[i] = np.r_[xy_list, np.array([[1, 1]])] #xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) if draw_gt_quad: draw.line([ tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([ tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([ tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]]) ], width=3, fill='yellow') if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) np.save(os.path.join(train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format(o_img_fname, d_wight, d_height)) train_img_list = os.listdir(train_image_dir) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) print('found %d train labels.' % len(train_label_list)) random.shuffle(train_val_set) val_count = int(cfg.validation_split_ratio * len(train_val_set)) with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: f_val.writelines(train_val_set[:val_count]) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: f_train.writelines(train_val_set[val_count:])
def preprocess(): data_dir = cfg.data_dir origin_image_dir = os.path.join(data_dir, cfg.origin_image_dir_name) origin_txt_dir = os.path.join(data_dir, cfg.origin_txt_dir_name) train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) # 如果训练图片文件夹不存在,创建他 if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) draw_gt_quad = cfg.draw_gt_quad show_gt_image_dir = os.path.join(data_dir, cfg.show_gt_image_dir_name) # 创建显示gt的图片文件夹 if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) # 创建显示act图片文件夹 show_act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) # 列出origin文件夹下的所有图片 o_img_list = os.listdir(origin_image_dir) print('found %d origin images.' % len(o_img_list)) train_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # d_wight, d_height = resize_image(im) # 根据cfg中设置的图片尺寸,求图片缩放比率 d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.size[0] scale_ratio_h = d_height / im.size[1] im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') show_gt_im = im.copy() # 在图片上绘制直线,这里先创建一个Draw对象 draw = ImageDraw.Draw(show_gt_im) with open(os.path.join(origin_txt_dir, o_img_fname[:-4] + '.txt'), 'r') as f: anno_list = f.readlines() # anno_list中的元素逐个输出时的结果 # 48.45,231.83,17.87,178.79,179.84,11.1,228.79,47.95,时尚袋袋 # 228.97,527.45,228.97,576.4,479.64,575.4,473.64,524.45,百事可乐 xy_list_array = np.zeros((len(anno_list), 4, 2)) # xy_list_array存储的是一张图片所有box的四个vertex坐标 # 读取一张图片的gt文本文件,在图片绘制出 for anno, i in zip(anno_list, range(len(anno_list))): anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) # xy_list = [ # [48.45,231.83], # [17.87,178.79], # [179.84,11.1], # [228.79,47.95] # ] xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w # 缩放了,所以w坐标位置发生了变化 xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h # 缩放了,所以h坐标位置发生了变化 xy_list = reorder_vertexes(xy_list) # 为四边形顶点重新排序,调整后的顺序为左上、左下、右下、右上 xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) # 如果是一个正立矩形,长宽内缩尺寸均为原矩形短边长度的0.2 shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) # shrink_1缩了x方向,y方向仍是原来的 if draw_gt_quad: draw.line([tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]])], width=3, fill='yellow') if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) # 保存仅仅进行了缩放的图片到指定的训练数据保存文件夹 np.save(os.path.join( train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) # 保存重排序(因为进行了缩放啊)后的四个顶点坐标们 if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) # 把绘制了线条的图保存到指定的文件夹 train_val_set.append('{},{},{}\n'.format(o_img_fname, # 训练数据又添一员,先保存到内存中,待最后统一写入到一个文件中啊 d_wight, d_height)) train_img_list = os.listdir(train_image_dir) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) print('found %d train labels.' % len(train_label_list)) random.shuffle(train_val_set) # 打乱所有的训练样本 val_count = int(cfg.validation_split_ratio * len(train_val_set)) # 按照训练、验证比例求得序号 with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: # 序号前的为验证样本 f_val.writelines(train_val_set[:val_count]) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: # 序号后的训练样本 f_train.writelines(train_val_set[val_count:])
def preprocess(): # calculate number of text, arrow, nock num_text = 0 num_arrow = 0 num_nock = 0 nock_samples = [] data_dir = cfg.data_dir origin_image_dir = os.path.join(data_dir, cfg.origin_image_dir_name) origin_txt_dir = os.path.join(data_dir, cfg.origin_txt_dir_name) train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) show_gt_image_dir = os.path.join(data_dir, cfg.show_gt_image_dir_name) if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) show_act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) o_img_list = os.listdir(origin_image_dir) print('found %d origin images.' % len(o_img_list)) train_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # d_wight, d_height = resize_image(im) d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') show_gt_im = im.copy() # draw on the img draw = ImageDraw.Draw(show_gt_im) if o_img_fname[:3] != 'cin': preffix, suffix = o_img_fname.split('_') suffix = str(suffix).replace('png', 'txt') with open( os.path.join(origin_txt_dir, preffix + '_label_' + suffix), 'r') as f: # load corresponding label file anno_list = f.readlines() xy_list_array = np.zeros((len(anno_list), 5, 2)) for anno, i in zip(anno_list, range(len(anno_list))): category, anno_colums = anno.strip().split('\t') anno_colums = anno_colums.strip().split(',') anno_array = np.array(anno_colums) xy_list = anno_array.astype(float).reshape((4, 2)) xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h xy_list = reorder_vertexes(xy_list) # set the class label, needed one-hot encoding if category == 'text': num_text = num_text + 1 xy_list_array[i] = np.r_[xy_list, np.array([[0, 0]])] # xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink( xy_list, cfg.shrink_side_ratio) elif category == 'nock': num_nock = num_nock + 1 xy_list_array[i] = np.r_[xy_list, np.array([[0, 1]])] _, shrink_xy_list, _ = shrink(xy_list, 0) shrink_1, _, long_edge = shrink(xy_list, 0.5) # nock_samples.append(current_file) elif category == 'arrow': num_arrow = num_arrow + 1 xy_list_array[i] = np.r_[xy_list, np.array([[1, 0]])] _, shrink_xy_list, _ = shrink(xy_list, 0) shrink_1, _, long_edge = shrink(xy_list, 0.5) # else: # #category == 'predict' # xy_list_array[i] = np.r_[xy_list, np.array([[1, 1]])] # _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) # shrink_1, _, long_edge = shrink(xy_list, # cfg.shrink_side_ratio) if cfg.DEBUG: draw.line([ tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([ tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([ tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]]) ], width=3, fill='yellow') if cfg.DEBUG: im.save(os.path.join(train_image_dir, o_img_fname)) np.save( os.path.join(train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) if cfg.DEBUG: show_gt_im.save( os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format( o_img_fname, d_wight, d_height)) else: # load corresponding json label file category_list, coords_list, current_file = load_json_label( os.path.join(origin_txt_dir, o_img_fname[:-4] + '.json')) xy_list_array = np.zeros((len(category_list), 5, 2)) for category, coord, i in zip(category_list, coords_list, range(len(category_list))): xy_list = coord.astype(float) xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h xy_list = reorder_vertexes(xy_list) # set the class label, needed one-hot encoding if category == 'text': num_text = num_text + 1 xy_list_array[i] = np.r_[xy_list, np.array([[0, 0]])] # xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink( xy_list, cfg.shrink_side_ratio) elif category == 'nock': num_nock = num_nock + 1 xy_list_array[i] = np.r_[xy_list, np.array([[0, 1]])] _, shrink_xy_list, _ = shrink(xy_list, 0) shrink_1, _, long_edge = shrink(xy_list, 0.5) # nock_samples.append(current_file) elif category == 'arrow': num_arrow = num_arrow + 1 xy_list_array[i] = np.r_[xy_list, np.array([[1, 0]])] _, shrink_xy_list, _ = shrink(xy_list, 0) shrink_1, _, long_edge = shrink(xy_list, 0.5) # else: # #category == 'predict' # xy_list_array[i] = np.r_[xy_list, np.array([[1, 1]])] # _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) # shrink_1, _, long_edge = shrink(xy_list, # cfg.shrink_side_ratio) if cfg.DEBUG: draw.line([ tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([ tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([ tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]]) ], width=3, fill='yellow') if cfg.DEBUG: im.save(os.path.join(train_image_dir, o_img_fname)) np.save( os.path.join(train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) if cfg.DEBUG: show_gt_im.save( os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format( o_img_fname, d_wight, d_height)) # del category_list, coords_list train_img_list = os.listdir(train_image_dir) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) print('found %d train labels.' % len(train_label_list)) random.shuffle(train_val_set) val_count = int(cfg.validation_split_ratio * len(train_val_set)) with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: f_val.writelines(train_val_set[:val_count]) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: f_train.writelines(train_val_set[val_count:]) return num_text, num_nock, num_arrow, nock_samples
def preprocess(): dataset_dir = config.DATASET_DIR # 存放原来的图片 origin_image_dir = os.path.join(dataset_dir, config.ORIGIN_IMAGE_DIR_NAME) # 存放原来的 txt 标注 origin_label_dir = os.path.join(dataset_dir, config.ORIGIN_LABEL_DIR_NAME) if config.SAVE_RESIZED_IMAGE: # 存放 resize 后的图片 resized_image_dir = os.path.join(dataset_dir, config.RESIZED_IMAGE_DIR_NAME) if not os.path.exists(resized_image_dir): logger.info( 'resized_image_dir:{} does not exist, then creating'.format( resized_image_dir)) os.mkdir(resized_image_dir) else: logger.info( 'resized_image_dir:{} already exists, then deleting and creating' .format(resized_image_dir)) shutil.rmtree(resized_image_dir) os.mkdir(resized_image_dir) if config.SAVE_RESIZED_LABEL: # 存放 resize 后的标注 resized_label_dir = os.path.join(dataset_dir, config.RESIZED_LABEL_DIR_NAME) if not os.path.exists(resized_label_dir): logger.info( 'resized_label_dir:{} does not exist, then creating'.format( resized_label_dir)) os.mkdir(resized_label_dir) else: logger.info( 'resized_label_dir:{} already exists, then deleting and creating' .format(resized_label_dir)) shutil.rmtree(resized_label_dir) os.mkdir(resized_label_dir) if config.SAVE_DRAW_GT_QUAD_IMAGE: # 存放画有 gt quad 的 resized_image draw_gt_quad_image_dir = os.path.join( dataset_dir, config.DRAW_GT_QUAD_IMAGE_DIR_NAME) if not os.path.exists(draw_gt_quad_image_dir): logger.info( 'draw_gt_quad_image_dir:{} does not exist, then creating'. format(draw_gt_quad_image_dir)) os.mkdir(draw_gt_quad_image_dir) else: logger.info( 'draw_gt_quad_image_dir:{} already exists, then deleting and creating' .format(draw_gt_quad_image_dir)) shutil.rmtree(draw_gt_quad_image_dir) os.mkdir(draw_gt_quad_image_dir) origin_image_filenames = os.listdir(origin_image_dir) num_origin_images = len(origin_image_filenames) logger.info('Found {} origin images.'.format(num_origin_images)) for origin_image_filename, _ in zip(origin_image_filenames, tqdm(range(num_origin_images))): origin_image_filepath = os.path.join(origin_image_dir, origin_image_filename) logger.debug('Handling {} starts'.format(origin_image_filepath)) origin_image = cv2.imread(origin_image_filepath) # cv2.namedWindow('origin_image', cv2.WINDOW_NORMAL) # cv2.imshow('origin_image', origin_image) # cv2.waitKey(0) if origin_image is None: logger.warning('Reading {} failed'.format(origin_image_filepath)) continue image, window, scale, padding, crop = resize_image( origin_image, min_dim=config.IMAGE_MIN_DIM, max_dim=config.IMAGE_MAX_DIM, mode=config.IMAGE_RESIZE_MODE) origin_label_filename = origin_image_filename[:-4] + '.txt' origin_label_filepath = os.path.join(origin_label_dir, origin_label_filename) if osp.exists(origin_label_filepath): with open(origin_label_filepath, 'r') as f: annotation_list = f.readlines() else: logger.warning('{} does not exist'.format(origin_label_filepath)) continue # draw ground truth quad on image draw_gt_quad_image = image.copy() num_annotations = len(annotation_list) xy_list_array = np.zeros((num_annotations, 4, 2)) for idx, annotation in enumerate(annotation_list): annotation_columns = annotation.strip().split(',') annotation_array = np.array(annotation_columns) xy_list = np.reshape(annotation_array[:8].astype(np.float64), (4, 2)) xy_list = xy_list * scale # 所有 x 加上 left_padding xy_list[:, 0] += padding[1][0] # 所有 y 加上 top_padding xy_list[:, 1] += padding[0][0] # 希望得到从左上方的点开始逆时针旋转得到的所有点的坐标 xy_list = sort_xylist(xy_list) xy_list_array[idx] = xy_list _, shrink_xy_list, _ = shrink(xy_list, config.SHRINK_RATIO) shrink_long_edges_xy_list, _, first_long_edge_idx = shrink( xy_list, config.SHRINK_SIDE_RATIO) if config.DRAW_GT_QUAD: cv2.polylines(draw_gt_quad_image, [xy_list.astype('int').reshape(-1, 1, 2)], True, (0, 255, 0), 2) cv2.polylines(draw_gt_quad_image, [shrink_xy_list.astype('int').reshape(-1, 1, 2)], True, (255, 0, 0), 2) # 左右长上下短的情况 # ---------------- # |//////////////| # ---------------- # | | # | | # | | # ---------------- # |//////////////| # ---------------- if first_long_edge_idx == 0: # 上端阴影 cv2.polylines(draw_gt_quad_image, [ np.array([ xy_list[0], shrink_long_edges_xy_list[0], shrink_long_edges_xy_list[3], xy_list[3] ]).astype('int').reshape(-1, 1, 2) ], True, (255, 255, 0), 2) # 下端阴影 cv2.polylines(draw_gt_quad_image, [ np.array([ shrink_long_edges_xy_list[1], xy_list[1], xy_list[2], shrink_long_edges_xy_list[2] ]).astype('int').reshape(-1, 1, 2) ], True, (255, 255, 0), 2) # 左右短上下长的情况 # -------------------------- # |///| |///| # |///| |///| # |///| |///| # |///| |///| # --------------------------- elif first_long_edge_idx == 1: # 左端阴影 cv2.polylines(draw_gt_quad_image, [ np.array([ xy_list[0], xy_list[1], shrink_long_edges_xy_list[1], shrink_long_edges_xy_list[0] ]).astype('int').reshape(-1, 1, 2) ], True, (255, 255, 0), 2) # 右端阴影 cv2.polylines(draw_gt_quad_image, [ np.array([ shrink_long_edges_xy_list[3], shrink_long_edges_xy_list[2], xy_list[2], xy_list[3] ]).astype('int').reshape(-1, 1, 2) ], True, (255, 255, 0), 2) if config.SAVE_RESIZED_IMAGE: resized_image_filename = origin_image_filename resized_image_filepath = os.path.join(resized_image_dir, resized_image_filename) # cv2.namedWindow('resized_image', cv2.WINDOW_NORMAL) # cv2.imshow('resized_image', image) # cv2.waitKey(0) cv2.imwrite(resized_image_filepath, image) if config.SAVE_RESIZED_LABEL: resized_label_filename = origin_label_filename[:-4] + '.npy' resized_label_filepath = os.path.join(resized_label_dir, resized_label_filename) np.save(resized_label_filepath, xy_list_array) if config.SAVE_DRAW_GT_QUAD_IMAGE: draw_gt_quad_image_filepath = os.path.join(draw_gt_quad_image_dir, origin_image_filename) # cv2.namedWindow('draw_gt_quad_image', cv2.WINDOW_NORMAL) # cv2.imshow('draw_gt_quad_image', draw_gt_quad_image) # cv2.waitKey(0) cv2.imwrite(draw_gt_quad_image_filepath, draw_gt_quad_image) logger.debug('Handling {} ends'.format(origin_image_filepath)) if config.SAVE_TRAIN_VAL_IMAGE: resized_image_dir = osp.join(dataset_dir, config.RESIZED_IMAGE_DIR_NAME) resized_label_dir = osp.join(dataset_dir, config.RESIZED_LABEL_DIR_NAME) for dir_type, dir_path in zip( ['resized_image_dir', 'resized_label_dir'], [resized_image_dir, resized_label_dir]): if not osp.exists(dir_path) or len(os.listdir(dir_path)) == 0: logger.error('{}:{} does not exist or is empty'.format( dir_type, dir_path)) return if not every_image_has_corresponding_label(resized_image_dir, resized_label_dir): logger.error( 'Files in resized_image_dir does not match files in resized_label_dir' ) return resized_image_filenames = os.listdir(resized_image_dir) num_resized_images = len(resized_image_filenames) # 存放用于 train 的 image train_image_dir = os.path.join(dataset_dir, config.TRAIN_IMAGE_DIR_NAME) # 存放用于 validation 的 image val_image_dir = os.path.join(dataset_dir, config.VAL_IMAGE_DIR_NAME) for dir_type, dir_path in zip(['train_image_dir', 'val_image_dir'], [train_image_dir, val_image_dir]): if not os.path.exists(dir_path): logger.info('{}:{} does not exist, then creating'.format( dir_type, dir_path)) os.mkdir(dir_path) else: logger.info( '{}:{} already exists, then deleting and creating'.format( dir_type, dir_path)) shutil.rmtree(dir_path) os.mkdir(dir_path) num_val_images = int(config.VAL_SPLIT_RATIO * num_resized_images) logger.debug('Expecting {} val images'.format(num_val_images)) num_train_images = num_resized_images - num_val_images logger.debug('Expecting {} train images'.format(num_train_images)) for image_filename in resized_image_filenames[:num_train_images]: train_image_filepath = osp.join(train_image_dir, image_filename) resized_image_filepath = osp.join(resized_image_dir, image_filename) shutil.copy(resized_image_filepath, train_image_filepath) for image_filename in resized_image_filenames[num_train_images:]: val_image_filepath = osp.join(val_image_dir, image_filename) resized_image_filepath = osp.join(resized_image_dir, image_filename) shutil.copy(resized_image_filepath, val_image_filepath) logger.info('Found {} train images.'.format( len(os.listdir(train_image_dir)))) logger.info('Found {} val images.'.format( len(os.listdir(val_image_dir))))
def preprocess(): data_dir = cfg.data_dir origin_image_dir = os.path.join( data_dir, cfg.origin_image_dir_name) # 'image_10000/' origin_txt_dir = os.path.join(data_dir, cfg.origin_txt_dir_name) # 'txt_10000/' train_image_dir = os.path.join( data_dir, cfg.train_image_dir_name) # 'images_%s/' % train_task_id train_label_dir = os.path.join( data_dir, cfg.train_label_dir_name) # 'labels_%s/' % train_task_id if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) draw_gt_quad = cfg.draw_gt_quad # True show_gt_image_dir = os.path.join( data_dir, cfg.show_gt_image_dir_name) # 'show_gt_images_%s/' % train_task_id if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) show_act_image_dir = os.path.join( cfg.data_dir, cfg.show_act_image_dir_name) # 'show_act_images_%s/' % train_task_id if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) o_img_list = os.listdir(origin_image_dir) print('found %d origin images.' % len(o_img_list)) train_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # 打开每张图片 # d_wight, d_height = resize_image(im) d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') # 图片缩放 show_gt_im = im.copy() # draw on the img draw = ImageDraw.Draw(show_gt_im) with open(os.path.join(origin_txt_dir, o_img_fname[:-4] + '.txt'), 'r', encoding='UTF-8') as f: anno_list = f.readlines() xy_list_array = np.zeros((len(anno_list), 4, 2)) for anno, i in zip(anno_list, range(len(anno_list))): anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w # 坐标缩放 xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h xy_list = reorder_vertexes(xy_list) # 坐标顺序转换为统一格式 xy_list_array[i] = xy_list # 将groundtruth文本框内缩,论文为0.3,实际中发现太大,改为0.2 _, shrink_xy_list, _ = shrink( xy_list, cfg.shrink_ratio) # shrink_ratio=0.2,返回长短边都缩后的结果 shrink_1, _, long_edge = shrink( xy_list, cfg.shrink_side_ratio ) # shrink_side_ratio=0.6,返回仅长边收缩的结果以及长边下标 if draw_gt_quad: draw.line([ tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([ tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): # 框出头跟尾巴的像素 draw.line([ tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]]) ], width=3, fill='yellow') if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) np.save(os.path.join(train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) # 保存顺序一致处理后的坐标点集 if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format(o_img_fname, d_wight, d_height)) train_img_list = os.listdir(train_image_dir) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) print('found %d train labels.' % len(train_label_list)) random.shuffle(train_val_set) val_count = int(cfg.validation_split_ratio * len(train_val_set)) with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: f_val.writelines(train_val_set[:val_count]) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: f_train.writelines(train_val_set[val_count:])
def preprocess(): # 主目录,默认是icpr data_dir = cfg.data_dir # 原图片和标注目录 origin_image_dir = os.path.join(data_dir, cfg.origin_image_dir_name) origin_txt_dir = os.path.join(data_dir, cfg.origin_txt_dir_name) # 训练图片和标注目录 train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) # 是否画出边界 draw_gt_quad = cfg.draw_gt_quad # 画出边界的图保存的目录 show_gt_image_dir = os.path.join(data_dir, cfg.show_gt_image_dir_name) if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) show_act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) # 获取原图片目录下所有的文件名,组装成列表 o_img_list = os.listdir(origin_image_dir) print('found %d origin images.' % len(o_img_list)) train_val_set = [] # tqdm:python进度条扩展 for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): # 循环打开原图片目录下一个图片作为对象 with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # d_wight, d_height = resize_image(im) d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height # 将图片resize以适合训练要求,提高训练速度 im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') show_gt_im = im.copy() # draw on the img draw = ImageDraw.Draw(show_gt_im) # 获取对应图片的所有标注信息 with open(os.path.join(origin_txt_dir, o_img_fname[:-4] + '.txt'), 'r', encoding='utf-8') as f: anno_list = f.readlines() xy_list_array = np.zeros((len(anno_list), 4, 2)) # 在图片上画出标注的信息 for anno, i in zip(anno_list, range(len(anno_list))): anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h xy_list = reorder_vertexes(xy_list) xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) if draw_gt_quad: draw.line([tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]])], width=3, fill='yellow') # 是否生成 resize 过后的图片(没用画出 GT 的) if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) # 生成 npy 的标注信息 np.save(os.path.join( train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) # 是否生成画出了标注的图片目录 if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format(o_img_fname, d_wight, d_height)) train_img_list = os.listdir(train_image_dir) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) print('found %d train labels.' % len(train_label_list)) # 打乱图片 random.shuffle(train_val_set) val_count = int(cfg.validation_split_ratio * len(train_val_set)) # 生成train and val txt with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: f_val.writelines(train_val_set[:val_count]) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: f_train.writelines(train_val_set[val_count:])
def preprocess(): if (cfg.train_val_split_control == True): preprocess_2() else: data_dir = cfg.data_dir origin_image_dir = os.path.join(cfg.origin_data_dir, cfg.origin_image_dir_name) origin_txt_dir = os.path.join(cfg.origin_data_dir, cfg.origin_txt_dir_name) train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) # 转换后的训练集目录 train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) # 转换后的训练集标签目录 if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) draw_gt_quad = cfg.draw_gt_quad show_gt_image_dir = os.path.join(data_dir, cfg.show_gt_image_dir_name) if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) show_act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) # 训练数据预处理(resize + 标签转为numpy) o_img_list = os.listdir(origin_image_dir) # 返回指定的文件夹包含的文件或文件夹的名字的列表 print('found %d train origin images.' % len(o_img_list)) train_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): if (o_img_fname[-4:] == ".txt"): continue with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # 打开图片 # d_wight, d_height = resize_image(im) d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') # resize show_gt_im = im.copy() # draw on the img draw = ImageDraw.Draw(show_gt_im) with open(os.path.join(origin_txt_dir, o_img_fname[:-4] + '.txt'), 'r') as f: # 打开标签文件 anno_list = f.readlines() xy_list_array = np.zeros((len(anno_list), 4, 2)) # 标签转化为 numpy 的格式 for anno, i in zip(anno_list, range(len(anno_list))): # 对该标签文件的每一个框 anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) # 取前8位 xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w # 横坐标 resize xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h # 纵坐标 resize xy_list = reorder_vertexes(xy_list) # 排序坐标 xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) if draw_gt_quad: # 画出GT框 draw.line([tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]])], width=3, fill='yellow') if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) np.save(os.path.join( train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format(o_img_fname, d_wight, d_height)) train_img_list = os.listdir(train_image_dir) # 找到训练图像的目录 (resize过后的) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) # 找到训练标签的目录 (resize过后的) print('found %d train labels.' % len(train_label_list)) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: # 训练集图像名写入 f_train.writelines(train_val_set[:]) # 测试数据预处理(resize + 标签转为numpy) origin_test_dir = os.path.join(cfg.origin_data_dir, cfg.test_origin_image_dir_name) origin_txt_test_dir = os.path.join(cfg.origin_data_dir, cfg.test_origin_txt_dir_name) o_img_list = os.listdir(origin_test_dir) print('found %d origin images.' % len(o_img_list)) test_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): if (o_img_fname[-4:] == ".txt"): continue with Image.open(os.path.join(origin_test_dir, o_img_fname)) as im: # d_wight, d_height = resize_image(im) d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') # resize show_gt_im = im.copy() # draw on the img draw = ImageDraw.Draw(show_gt_im) with open(os.path.join(origin_txt_dir, o_img_fname[:-4] + '.txt'), 'r') as f: # 打开标签文件 anno_list = f.readlines() xy_list_array = np.zeros((len(anno_list), 4, 2)) # 标签转化为numpy for anno, i in zip(anno_list, range(len(anno_list))): anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h xy_list = reorder_vertexes(xy_list) xy_list_array[i] = xy_list _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) if draw_gt_quad: draw.line([tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') draw.line([tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for q_th in range(2): draw.line([tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]])], width=3, fill='yellow') if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) np.save(os.path.join( train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) test_val_set.append('{},{},{}\n'.format(o_img_fname, d_wight, d_height)) train_img_list = os.listdir(train_image_dir) # 找到训练图像的目录 (resize过后的) print('found %d (train + test) images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) # 找到训练标签的目录 (resize过后的) print('found %d (train + test) labels.' % len(train_label_list)) # random.shuffle(train_val_set) # 打乱 # val_count = int(cfg.validation_split_ratio * len(train_val_set)) # 划分验证集和训练集 with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: # 将验证集图像名写入 f_val.writelines(test_val_set[:])
def preprocess(): # '/Users/xingoo/Documents/dataset/train_1000' data_dir = cfg.data_dir # 原始的图片文件 image_1000 origin_image_dir = os.path.join(data_dir, cfg.origin_image_dir_name) # 原始的label文件 txt_1000 格式为八个点的坐标和文本 origin_txt_dir = os.path.join(data_dir, cfg.origin_txt_dir_name) # 输出的文件夹 images_1000 train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) # 输出的label文件夹 labels_1000 train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) if not os.path.exists(train_image_dir): os.mkdir(train_image_dir) if not os.path.exists(train_label_dir): os.mkdir(train_label_dir) # 是否显示ground_truth的背景图片 draw_gt_quad = cfg.draw_gt_quad # 显示真实的边框图 show_gt_images_1000 show_gt_image_dir = os.path.join(data_dir, cfg.show_gt_image_dir_name) if not os.path.exists(show_gt_image_dir): os.mkdir(show_gt_image_dir) # 显示真实的像素点图 show_act_images_1000 show_act_image_dir = os.path.join(cfg.data_dir, cfg.show_act_image_dir_name) if not os.path.exists(show_act_image_dir): os.mkdir(show_act_image_dir) o_img_list = os.listdir(origin_image_dir) print('found %d origin images.' % len(o_img_list)) train_val_set = [] for o_img_fname, _ in zip(o_img_list, tqdm(range(len(o_img_list)))): with Image.open(os.path.join(origin_image_dir, o_img_fname)) as im: # d_wight, d_height = resize_image(im) # 获得原始尺寸,并计算缩放到标准的尺度的比例,目前是736*736 d_wight, d_height = cfg.max_train_img_size, cfg.max_train_img_size scale_ratio_w = d_wight / im.width scale_ratio_h = d_height / im.height im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB') show_gt_im = im.copy() # draw on the img # 读取标记文件 draw = ImageDraw.Draw(show_gt_im) with open(os.path.join(origin_txt_dir, o_img_fname[:-4] + '.txt'), 'r', encoding='utf-8') as f: anno_list = f.readlines() # 创建(n,4,2)的大小 xy_list_array = np.zeros((len(anno_list), 4, 2)) # for anno, i in zip(anno_list, range(len(anno_list))): anno_colums = anno.strip().split(',') anno_array = np.array(anno_colums) xy_list = np.reshape(anno_array[:8].astype(float), (4, 2)) # 尺度缩放 xy_list[:, 0] = xy_list[:, 0] * scale_ratio_w xy_list[:, 1] = xy_list[:, 1] * scale_ratio_h # xy_list为缩放后的8个点的值,按照顺序对8个位置进行排序 xy_list = reorder_vertexes(xy_list) xy_list_array[i] = xy_list # todo 这里是对图像做了增广吗? _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) if draw_gt_quad: # 原始框 draw.line([ tuple(xy_list[0]), tuple(xy_list[1]), tuple(xy_list[2]), tuple(xy_list[3]), tuple(xy_list[0]) ], width=2, fill='green') # 缩小框 draw.line([ tuple(shrink_xy_list[0]), tuple(shrink_xy_list[1]), tuple(shrink_xy_list[2]), tuple(shrink_xy_list[3]), tuple(shrink_xy_list[0]) ], width=2, fill='blue') vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] # 头和尾的数据框 for q_th in range(2): draw.line([ tuple(xy_list[vs[long_edge][q_th][0]]), tuple(shrink_1[vs[long_edge][q_th][1]]), tuple(shrink_1[vs[long_edge][q_th][2]]), tuple(xy_list[vs[long_edge][q_th][3]]), tuple(xy_list[vs[long_edge][q_th][4]]) ], width=3, fill='yellow') if cfg.gen_origin_img: im.save(os.path.join(train_image_dir, o_img_fname)) np.save(os.path.join(train_label_dir, o_img_fname[:-4] + '.npy'), xy_list_array) if draw_gt_quad: show_gt_im.save(os.path.join(show_gt_image_dir, o_img_fname)) train_val_set.append('{},{},{}\n'.format(o_img_fname, d_wight, d_height)) train_img_list = os.listdir(train_image_dir) print('found %d train images.' % len(train_img_list)) train_label_list = os.listdir(train_label_dir) print('found %d train labels.' % len(train_label_list)) random.shuffle(train_val_set) # 10%作为验证集,90%作为测试集 val_count = int(cfg.validation_split_ratio * len(train_val_set)) with open(os.path.join(data_dir, cfg.val_fname), 'w') as f_val: f_val.writelines(train_val_set[:val_count]) with open(os.path.join(data_dir, cfg.train_fname), 'w') as f_train: f_train.writelines(train_val_set[val_count:])