CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutSign UpSign In
hukaixuan19970627

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

GitHub Repository: hukaixuan19970627/yolov5_obb
Path: blob/master/DOTA_devkit/prepare_dota1_ms.py
Views: 475
1
import os
2
import os.path as osp
3
4
from DOTA_devkit.DOTA2JSON import generate_json_labels
5
from DOTA_devkit.DOTA2COCO_poly import DOTA2COCOTrain, DOTA2COCOTest, wordname_15
6
7
from DOTA_devkit.ImgSplit_multi_process import splitbase as splitbase_trainval
8
from DOTA_devkit.SplitOnlyImage_multi_process import \
9
splitbase as splitbase_test
10
11
12
def mkdir_if_not_exists(path):
13
if not osp.exists(path):
14
os.mkdir(path)
15
16
def prepare_multi_scale_data(src_path, dst_path, gap=200, subsize=1024, scales=[0.5, 1.0, 1.5], num_process=32):
17
"""Prepare DOTA split data and labels
18
Args:
19
src_path: dataset path
20
dst_path: output path
21
gap: overlapping area
22
subsize: size of chip image
23
scales: multi-scale settings
24
num_process: num of processer
25
"""
26
dst_train_path = osp.join(dst_path, 'train_split')
27
dst_val_path = osp.join(dst_path, 'val_split')
28
dst_trainval_path = osp.join(dst_path, 'trainval_split')
29
dst_test_base_path = osp.join(dst_path, 'test_split')
30
dst_test_path = osp.join(dst_path, 'test_split/images')
31
# make dst path if not exist
32
mkdir_if_not_exists(dst_path)
33
mkdir_if_not_exists(dst_train_path)
34
mkdir_if_not_exists(dst_val_path)
35
mkdir_if_not_exists(dst_test_base_path)
36
mkdir_if_not_exists(dst_test_path)
37
# split train data
38
print('split train data')
39
split_train = splitbase_trainval(osp.join(src_path, 'train'), dst_train_path,
40
gap=gap, subsize=subsize, num_process=num_process)
41
for scale in scales:
42
split_train.splitdata(scale)
43
print('split val data')
44
# split val data
45
split_val = splitbase_trainval(osp.join(src_path, 'val'), dst_val_path,
46
gap=gap, subsize=subsize, num_process=num_process)
47
for scale in scales:
48
split_val.splitdata(scale)
49
# split test data
50
print('split test data')
51
split_test = splitbase_test(osp.join(src_path, 'test/images'), dst_test_path,
52
gap=gap, subsize=subsize, num_process=num_process)
53
for scale in scales:
54
split_test.splitdata(scale)
55
56
# prepare trainval data
57
print('move train val to trainval')
58
mkdir_if_not_exists(dst_trainval_path)
59
os.system(
60
'mv {}/* {}'.format(dst_train_path, dst_trainval_path))
61
os.system('find '+dst_val_path+'/images/ -name "*.png" -exec mv {} ' +
62
dst_trainval_path + '/images/ \\;')
63
os.system('find '+dst_val_path+'/labelTxt/ -name "*.txt" -exec mv {} ' +
64
dst_trainval_path + '/labelTxt/ \\;')
65
66
print('generate labels with json format')
67
generate_json_labels(dst_trainval_path, osp.join(
68
dst_trainval_path, 'trainval.json'))
69
generate_json_labels(dst_test_base_path, osp.join(
70
dst_test_base_path, 'test.json'), trainval=False)
71
print('generate labels with coco format')
72
DOTA2COCOTrain(dst_trainval_path,
73
osp.join(dst_trainval_path, 'trainval_coco_8point.json'),
74
wordname_15)
75
DOTA2COCOTest(dst_test_base_path,
76
osp.join(dst_test_base_path, 'test_coco_8point.json'),
77
wordname_15)
78
79
80
if __name__ == '__main__':
81
# single scale
82
prepare_multi_scale_data('/data1/dataset_demo/DOTA_demo/',
83
'/data1/OrientedRepPoints/data/dota_1024', scales=[1.0], gap=200)
84
# multi scale
85
# prepare_multi_scale_data('/mnt/SSD/lwt_workdir/data/dota_new/',
86
# '/mnt/SSD/lwt_workdir/data/dota_1024_ms', scales=[0.5, 1.0, 1.5], gap=500)
87
print('done')
88
89