DMNet复现(一)之数据准备篇：Density map guided object detection in aerial image

一、生成密度图

密度图标签生成

采用以下代码，生成训练集密度图gt：

import cv2
import glob
import h5py
import scipy
import pickle
import numpy as np
from PIL import Image
from itertools import islice
from tqdm import tqdm
from matplotlib import pyplot as plt
from sortedcontainers import SortedDict
from scipy.ndimage.filters import gaussian_filter
from scipy.spatial import KDTree
import argparse"""
Code for DMnet, density map ground truth generation
Author: Changlin Li
Code revised on : 7/15/2020Given dataset(train/val/test) generate ground truth for given dataset.
Default format for source data: The input images are in jpg format and raw annotations are in txt format 
(Based on Visiondrone 2018/19/20 dataset)Sample code to run:python Generate_density_map_official.py . gaussian_kernels.pkl distances_dict.pkl --mode val
"""# point_class_pair = {}
# annotation_stats = {0: 17, 1: 14, 2: 20, 3: 32, 4: 35, 5: 45, 6: 29, 7: 30, 8: 46, 9: 18}
# min_sigma, max_sigma = min(annotation_stats.values()), max(annotation_stats.values())
# print(min_sigma, max_sigma)def get_img_paths(path_sets):"""Return all images from all pathes in 'path_sets'"""img_paths = []for path in path_sets:for img_path in glob.glob(os.path.join(path, '*.jpg')):img_paths.append(img_path)return img_pathsdef save_computed_density(density_map, out_path):"""Save density map to h5py format"""with h5py.File(out_path, 'w') as hf:hf['density'] = density_mapdef compute_sigma(gt_count, distance=None, min_sigma=1, method=1, fixed_sigma=15):"""Compute sigma for gaussian kernel with different methods :* method = 1 : sigma = (sum of distance to 3 nearest neighbors) / 10* method = 2 : sigma = distance to nearest neighbor* method = 3 : sigma = fixed value** if sigma lower than threshold 'min_sigma', then 'min_sigma' will be used** in case of one point on the image sigma = 'fixed_sigma'"""if gt_count > 1 and distance is not None:if method == 1:sigma = np.mean(distance[1:4]) * 0.1elif method == 2:sigma = distance[1]elif method == 3:sigma = fixed_sigmaelse:sigma = fixed_sigmaif sigma < min_sigma:sigma = min_sigmareturn sigmadef find_closest_key(sorted_dict, key):"""Find closest key in sorted_dict to 'key'"""keys = list(islice(sorted_dict.irange(minimum=key), 1))keys.extend(islice(sorted_dict.irange(maximum=key, reverse=True), 1))return min(keys, key=lambda k: abs(key - k))def gaussian_filter_density(non_zero_points, map_h, map_w, distances=None, kernels_dict=None, min_sigma=2, method=1,const_sigma=15):"""Fast gaussian filter implementation : using precomputed distances and kernels"""gt_count = non_zero_points.shape[0]density_map = np.zeros((map_h, map_w), dtype=np.float32)for i in range(gt_count):point_x, point_y, category = non_zero_points[i]sigma = compute_sigma(gt_count, distances[i], min_sigma=min_sigma, method=method, fixed_sigma=const_sigma)# closest_sigma = annotation_stats[category]closest_sigma = find_closest_key(kernels_dict, sigma)# print(i,closest_sigma)kernel = kernels_dict[closest_sigma]full_kernel_size = kernel.shape[0]kernel_size = full_kernel_size // 2min_img_x = max(0, point_x - kernel_size)min_img_y = max(0, point_y - kernel_size)max_img_x = min(point_x + kernel_size + 1, map_w - 1)max_img_y = min(point_y + kernel_size + 1, map_h - 1)assert max_img_x > min_img_xassert max_img_y > min_img_ykernel_x_min = kernel_size - point_x if point_x <= kernel_size else 0kernel_y_min = kernel_size - point_y if point_y <= kernel_size else 0kernel_x_max = kernel_x_min + max_img_x - min_img_xkernel_y_max = kernel_y_min + max_img_y - min_img_yassert kernel_x_max > kernel_x_minassert kernel_y_max > kernel_y_mindensity_map[min_img_y:max_img_y, min_img_x:max_img_x] += kernel[kernel_y_min:kernel_y_max,kernel_x_min:kernel_x_max]return density_mapdef get_gt_dots(ann_path, img_height, img_width, mode="train"):"""Load Matlab file with ground truth labels and save it to numpy array.** cliping is needed to prevent going out of the array"""txt_list = open(ann_path, 'r').readlines()gt = format_label(mode, txt_list)assert gt.shape[1] == 3gt[:, 0] = gt[:, 0].clip(0, img_width - 1)gt[:, 1] = gt[:, 1].clip(0, img_height - 1)return gtdef set_circles_on_img(image, bbox_list, circle_size=2):"""Set circles on images at centers of bboxes in bbox_list"""for bbox in bbox_list:cv2.circle(image, (bbox[0], bbox[1]), circle_size, (255, 0, 0), -1)return imagedef generate_gaussian_kernels(out_kernels_path='gaussian_kernels.pkl', round_decimals=3, sigma_threshold=4, sigma_min=0,sigma_max=20, num_sigmas=801):"""Computing gaussian filter kernel for sigmas in linspace(sigma_min, sigma_max, num_sigmas) and saving them to dict.     """if os.path.exists(out_kernels_path):# If kernel has been pre-computed, then returnprint("Kernel already created!\nExiting...\n")returnkernels_dict = dict()sigma_space = np.linspace(sigma_min, sigma_max, num_sigmas)for sigma in tqdm(sigma_space):sigma = np.round(sigma, decimals=round_decimals)kernel_size = np.ceil(sigma * sigma_threshold).astype(np.int)img_shape = (kernel_size * 2 + 1, kernel_size * 2 + 1)img_center = (img_shape[0] // 2, img_shape[1] // 2)arr = np.zeros(img_shape)arr[img_center] = 1arr = scipy.ndimage.filters.gaussian_filter(arr, sigma, mode='constant')kernel = arr / arr.sum()kernels_dict[sigma] = kernelprint(f'Computed {len(sigma_space)} gaussian kernels. Saving them to {out_kernels_path}')with open(out_kernels_path, 'wb') as f:pickle.dump(kernels_dict, f)def compute_distances(out_dist_path='distances_dict.pkl', raw_label_dir='D:/BaiduNetdiskDownload/VisDrone', n_neighbors=4,leafsize=1024, data_limit=None, mode="train", img_affix=".jpg"):if os.path.exists(out_dist_path):# If distance has been computed, then directly load distance file.print("Distrance pre-computation already created!\nExiting...\n")returndistances_dict = dict()full_img_paths = glob.glob(f'{raw_label_dir}/VisDrone2019-DET-train/images/*' + img_affix) + \glob.glob(f'{raw_label_dir}/VisDrone2019-DET-val/images/*' + img_affix) + \glob.glob(f'{raw_label_dir}/VisDrone2019-DET-test-dev/images/*' + img_affix)if data_limit and data_limit < len(full_img_paths):full_img_paths = full_img_paths[:data_limit]for img_path in tqdm(full_img_paths):ann_path = img_path.replace(img_affix, '.txt')ann_path = ann_path.replace("images", "annotations")img = plt.imread(img_path)non_zero_points = get_gt_dots(ann_path, *img.shape[0:2], mode=mode)tree = KDTree(non_zero_points.copy(), leafsize=leafsize)  # build kdtreedistances, _ = tree.query(non_zero_points, k=n_neighbors)  # query kdtreedistances_dict[img_path] = distancesprint(f'Distances computed for {len(full_img_paths)}. Saving them to {out_dist_path}')with open(out_dist_path, 'wb') as f:pickle.dump(distances_dict, f)def format_label(mode, txt_list):format_data = []# required format: xmin, ymin, xmax, ymax, class_id, clockwise direction# Given format: <bbox_left>,<bbox_top>,<bbox_width>,<bbox_height>,class_idfor idx, i in enumerate(txt_list):coord_raw = [int(x) for x in i.replace("\n", "").split(',') if len(x) != 0]coord = coord_raw[:6]# print(coord)if len(coord) != 6:# 4 coord + 1 classprint("Failed to parse annotation!")exit()# if coord[-1] not in class_list and coord[-1]>len(class_list):#     print('warning found a new label :', coord[-1])#     exit()if coord[2] <= 0 or coord[3] <= 0:print("Error encountered!\nFind out 0 height(width)!")print("This bounding box has been discarded! ")continue# print("Pull out corrd matrix:\n")# print(coord)# exit(-1)if not 0 < coord[-1] < 11:# class 0 and 11 are not in our interestcontinueif mode == "VisDrone2019-DET-val" or "VisDrone2019-DET-test":# in this case, score is the last 2 element.# No consideration for score 0 in evalif int(coord[-2]) == 0:continueif int(coord_raw[-2]) == 2:continuebbox_left, bbox_top = coord[0], coord[1]bbox_right, bbox_bottom = coord[0] + coord[2], coord[1] + coord[3]# Scale class number back to range 0-9center_x, center_y = int((bbox_left + bbox_right) * 0.5), int((bbox_top + bbox_bottom) * 0.5)format_data.append([center_x, center_y, coord[-1] - 1])# if not filename:#     continue# if filename not in point_class_pair:#     point_class_pair[filename] = {}# coord_pair = str(center_x) + " " + str(center_y)# if coord_pair not in point_class_pair[filename]:#     point_class_pair[filename][coord_pair] = coord[-1] - 1# else:#     if point_class_pair[filename][coord_pair] != coord[-1] - 1:#         assert True, \#             "duplicate coordination shows in current file : " + str(filename)return np.array(format_data)def parse_args():parser = argparse.ArgumentParser(description='DMNet--Density map ground truth generation')parser.add_argument('root_dir', default=".",help='the path for source data')parser.add_argument('precomputed_kernels_path', default="gaussian_kernels.pkl",help='the path to save precomputed kernels')parser.add_argument('precomputed_distances_path', default="distances_dict.pkl",help='the path to save precomputed distance')parser.add_argument('--image_prefix', default=".jpg", help='the path to save precomputed distance')parser.add_argument('--mode', default="train", help='Indicate if you are working on train/val/test set')parser.add_argument('--showden', action='store_true', help='show results')args = parser.parse_args()return argsif __name__ == "__main__":# General setupargs = parse_args()data_limit = Noneprecomputed_kernels_path = args.precomputed_kernels_pathprecomputed_distances_path = args.precomputed_distances_pathimg_affix = args.image_prefixshowden = args.showdenmode = args.moderoot_dir = args.root_dirmin_sigma = 0max_sigma = 20# create dir to save train/val density mapif not os.path.exists(os.path.join('/root/autodl-tmp/VisDrone2019/VisDrone2019-DET-train', 'dens')):os.makedirs(os.path.join('/root/autodl-tmp/VisDrone2019/VisDrone2019-DET-train', 'dens'), exist_ok=True)if not os.path.exists(os.path.join('/root/autodl-tmp/VisDrone2019/VisDrone2019-DET-val', 'dens')):os.makedirs(os.path.join('/root/autodl-tmp/VisDrone2019/VisDrone2019-DET-val', 'dens'), exist_ok=True)if not os.path.exists(os.path.join('/root/autodl-tmp/VisDrone2019/VisDrone2019-DET-test-dev', 'dens')):os.makedirs(os.path.join('/root/autodl-tmp/VisDrone2019/VisDrone2019-DET-test-dev', 'dens'), exist_ok=True)# create pre-computed kernel to speed up density map generationgenerate_gaussian_kernels(precomputed_kernels_path, round_decimals=3, sigma_threshold=4,sigma_min=min_sigma, sigma_max=max_sigma, num_sigmas=801)with open(precomputed_kernels_path, 'rb') as f:kernels_dict = pickle.load(f)kernels_dict = SortedDict(kernels_dict)# uncomment to generate and save dict with distancescompute_distances(out_dist_path=precomputed_distances_path, raw_label_dir=root_dir, mode=mode)with open(precomputed_distances_path, 'rb') as f:distances_dict = pickle.load(f)# print(distances_dict)data_root = modeimg_paths = glob.glob(f'{root_dir}/{data_root}/images/*.jpg')method = 3const_sigma = 15with open(str(mode) + ".txt", "w") as fileloader:# Prepared for the training algorithms that requires a txt output file# with all input images listedfor img_path in tqdm(img_paths):fileloader.write(img_path)fileloader.write("\n")data_folder, img_sub_path = img_path.split('images')ann_path = img_path.replace(img_affix, '.txt')ann_path = ann_path.replace("images", 'annotations')# load img and gtimg = Image.open(img_path)# print(img_path)width, height = img.sizegt_points = get_gt_dots(ann_path, height, width, mode=mode)distances = distances_dict[img_path]density_map = gaussian_filter_density(gt_points, height, width, distances,kernels_dict, min_sigma=min_sigma, method=method,const_sigma=const_sigma)den_name = os.path.join(root_dir, data_root, "dens", img_path.split("/")[-1].replace("jpg", "npy"))print(den_name)if showden:plt.imshow(img)plt.imshow(density_map, alpha=0.75)plt.show()else:np.save(den_name, density_map)

使用以下命令进行创建：

# 生成训练集gt密度图
python image_cropping/Generate_density_map_official.py /root/autodl-tmp/VisDrone2019 gaussian_kernels.pkl distances_dict.pkl --mode VisDrone2019-DET-train
# 生成验证集gt密度图
python image_cropping/Generate_density_map_official.py /root/autodl-tmp/VisDrone2019 gaussian_kernels.pkl distances_dict.pkl --mode VisDrone2019-DET-val
# 生成测试集gt密度图
python image_cropping/Generate_density_map_official.py /root/autodl-tmp/VisDrone2019 gaussian_kernels.pkl distances_dict.pkl --mode VisDrone2019-DET-test-dev

在以上文件夹 生成了如下的dens文件夹 保存了密度图的npy文件。

密度图预测

根据官方仓库说明 下载MCNN代码 进行训练预测即可

裁剪区域

# 生成训练集
python image_cropping/density_slide_window_official.py /root/autodl-tmp/VisDrone2019 70_70 0.08 --output_folder /root/autodl-tmp/VisDrone2019/Crop_70_0.08  --mode VisDrone2019-DET-train
# 生成验证集
python image_cropping/density_slide_window_official.py /root/autodl-tmp/VisDrone2019 70_70 0.08 --output_folder /root/autodl-tmp/VisDrone2019/Crop_70_0.08  --mode VisDrone2019-DET-val# 生成测试集
python image_cropping/density_slide_window_official.py /root/autodl-tmp/VisDrone2019 70_70 0.08 --output_folder /root/autodl-tmp/VisDrone2019/Crop_70_0.08  --mode VisDrone2019-DET-test-dev

部分切分后的效果图，可以看到部分目标还是存在截断现象，是一个优化点。





由于不清楚是否会将crop后的图像与标签，用于到训练过程中，这里我就不生成训练的crop图和标签了，下一步准备直接用训练好的模型进行融合测试。论文是采用验证集进行测试的，这儿我准备把验证集和测试集（test-dev）都测试一遍。

最后按以下结构需要组织一下文件

最后我的目录结构如下

数据格式转换

需要先把txt转为voc，再把voc转为coco

# 生成测试集的voc格式
python fusion_detection/create_VOC_annotation_official.py /root/autodl-tmp/VisDrone2019 --output_folder VisDrone2019_finally --mode VisDrone2019-DET-test-dev
# 生成测试集的coco格式
python fusion_detection/VOC2coco_official.py /root/autodl-tmp/VisDrone2019/VisDrone2019_finally --mode VisDrone2019-DET-test-dev

这过程太耗时间了。最后把中间文件夹删除，整个格式如以下：

其中coco的json文件在以下位置

过程太繁琐了。

有需要生成后的数据，可以联系！