def prepare_uavdt()

in scripts/data_prepare.py [0:0]

• 74 lines of code
• 9 McCabe index (conditional complexity)

def prepare_uavdt():
    root = opt.dataset
    data_dir = join(root, 'UAV-benchmark-M')
    attr_dir = join(root, 'M_attr')
    label_dir = join(root, 'UAV-benchmark-MOTD_v1.0', 'GT')
    split_dir = join(root, 'split')
    
    # mask images, it takes minutes
    for i, path in enumerate(glob(join(label_dir, '*ignore.txt'))):
        labels = np.loadtxt(path, usecols=(0, 2, 3, 4, 5, 8), dtype=int, delimiter=',')
        vid_name = basename(path).split('_')[0]
        for frameID, x1, y1, w, h, _ in tqdm(labels, desc='%02d/50' % (i + 1)):
            masked_path = join(data_dir, vid_name, 'img%06d_masked.jpg' % frameID)
            input_path = masked_path if exists(masked_path) else masked_path.replace('_masked.jpg', '.jpg')
            image = cv2.imread(input_path)
            image[y1:y1 + h, x1:x1 + w] = (127, 127, 127)
            cv2.imwrite(masked_path, image)

    data_split = {}
    for mode in ['train', 'test']:
        data_split[mode] = []
        for path in glob(join(attr_dir, mode, '*.txt')):
            data_split[mode].append(basename(path)[:5])
    k = 10
    sep = len(data_split['train']) // k
    random.shuffle(data_split['train'])
    data_split['train'], data_split['valid'] = data_split['train'][sep:], data_split['train'][:sep]

    os.mkdir(split_dir)
    for mode in ['train', 'valid', 'test']:
        with open(join(split_dir, '%s_video.txt' % mode), 'w+') as f:
            f.writelines(vid + '\n' for vid in data_split[mode])
        image_paths = []
        for video_name in tqdm(data_split[mode], desc=mode):
            ignore_path = join(label_dir, '%s_gt_ignore.txt' % video_name)
            label_path = join(label_dir, '%s_gt_whole.txt' % video_name)
        
            # the warning caused by empty file doesn't matter
            ignores = np.loadtxt(ignore_path, usecols=(0, 2, 3, 4, 5), dtype=int, delimiter=',')
            ignore_dict = {}
            for frameID, x1, y1, w, h in ignores:
                xyxy = np.array([[x1, y1, x1 + w, y1 + h]])
                if frameID in ignore_dict:
                    ignore_dict[frameID] = np.concatenate((ignore_dict[frameID], xyxy), axis=0)
                else:
                    ignore_dict[frameID] = xyxy
        
            labels = np.loadtxt(label_path, usecols=(0, 2, 3, 4, 5, 8), dtype=int, delimiter=',')
            label_dict = {}
            for frameID, x1, y1, w, h, cls in labels:
                xc, yc = x1 + w / 2., y1 + h / 2.
                if frameID in ignore_dict:
                    ignore_regions = ignore_dict[frameID]
                    if np.logical_and(
                            np.logical_and(ignore_regions[:, 0] < xc, xc < ignore_regions[:, 2]),
                            np.logical_and(ignore_regions[:, 1] < yc, yc < ignore_regions[:, 3])
                    ).sum() > 0:
                        continue
            
                box = [cls, xc, yc, w, h]
                if frameID in label_dict:
                    label_dict[frameID].append(box)
                else:
                    label_dict[frameID] = [box]
        
            for frameID, bboxes in label_dict.items():
                image_path = join(data_dir, video_name, 'img%06d_masked.jpg' % frameID)
                if not exists(image_path):
                    image_path = image_path.replace('_masked.jpg', '.jpg')
            
                image = cv2.imread(image_path)
                height, width, _ = image.shape
                label_path = image_path.replace('.jpg', '.txt')
                with open(label_path, 'w+') as f:
                    for cls, xc, yc, w, h in bboxes:
                        assert 1 <= cls <= 3  # cls - 1
                        # treat all categories to one, 'car'
                        f.write('%d %.6f %.6f %.6f %.6f\n' % (0, xc / width, yc / height, w / width, h / height))

                gen_mask(label_path, image)
                image_paths.append(image_path + '\n')
    
        with open(join(split_dir, '%s.txt' % mode), 'w+') as f:
            f.writelines(image_paths)
        if mode == 'train':
            with open(join(split_dir, '%s_ds.txt' % mode), 'w+') as f:
                f.writelines(image_paths[::10])