-
Notifications
You must be signed in to change notification settings - Fork 2
/
gen_labels_cmc.py
111 lines (92 loc) · 4.28 KB
/
gen_labels_cmc.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
'''
https://github.com/ifzhang/FairMOT
You can train FairMOT on custom dataset by following several steps bellow:
1. Generate one txt label file for one image. Each line of the txt label file represents one object. The format of the line is: "class id x_center/img_width y_center/img_height w/img_width h/img_height". You can modify src/gen_labels_16.py to generate label files for your custom dataset.
2. Generate files containing image paths. The example files are in src/data/. Some similar code can be found in src/gen_labels_crowd.py
3. Create a json file for your custom dataset in src/lib/cfg/. You need to specify the "root" and "train" keys in the json file. You can find some examples in src/lib/cfg/.
4. Add --data_cfg '../src/lib/cfg/your_dataset.json' when training.
'''
import glob
import os.path as osp
import os
import shutil
import numpy as np
from numpy.random import default_rng
def mkdirs(d):
if os.path.exists(d):
shutil.rmtree(d)
if not osp.exists(d):
os.makedirs(d)
def labels_with_ids():
seq_root = 'D:/dataset/tracking/CMC/data/images'
label_root = 'D:/dataset/tracking/CMC/data/labels_with_ids/'
seqs = [s for s in os.listdir(seq_root)]
mkdirs(label_root)
for idx, seq in enumerate(seqs):
for cam_num in [1, 2, 3, 4]:
# CMC dataset is a four-camera 1920x1024 resolution dataset
seq_width = 1920
seq_height = 1024
gt_txt = osp.join(seq_root, seq, "GT_" + seq + "_Cam" + str(cam_num) + ".txt")
gt = np.loadtxt(gt_txt, dtype=np.float64, delimiter=' ')
gt = gt[np.lexsort((gt[:, 0], gt[:, 1]))] # convert to MOT GT format
seq_label_root = osp.join(label_root, seq, "Cam_" + str(cam_num))
mkdirs(seq_label_root)
for fid, tid, x, y, w, h, _, _, _, _ in gt:
fid = int(fid)
tid_curr = int(tid) + 10 ** idx # use to distinguish IDs among dataset CMC1, CMC2, CMC3, CMC4, CMC5
x += w / 2
y += h / 2
label_fpath = osp.join(seq_label_root, '{:06d}.txt'.format(fid))
label_str = '0 {:d} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
tid_curr, x / seq_width, y / seq_height, w / seq_width, h / seq_height)
with open(label_fpath, 'a') as f:
f.write(label_str)
# End
# End
# End
def gen_image_paths():
seq_root = '../../data/images/'
seqs = [s for s in os.listdir(seq_root)]
if os.path.exists("cmc.train"):
os.remove("cmc.train")
for idx, seq in enumerate(seqs):
for cam_num in [1, 2, 3, 4]:
dataset_dir = osp.join(seq_root, seq, "Cam_" + str(cam_num) + "/*.png")
imgs = glob.glob(dataset_dir)
for img_dir in imgs:
with open("cmc.train", 'a') as f:
f.write(img_dir[len(seq_root):] + '\n')
def gen_image_paths_half():
seq_root = '../../data/images/'
seqs = [s for s in os.listdir(seq_root)]
if os.path.exists("cmc_half.train"):
os.remove("cmc_half.train")
rng = default_rng()
for idx, seq in enumerate(seqs):
for cam_num in [1, 2, 3, 4]:
dataset_dir = osp.join(seq_root, seq, "Cam_" + str(cam_num) + "/*.png")
imgs = glob.glob(dataset_dir)
half_ind = rng.choice(len(imgs), size=int(len(imgs) / 2), replace=False)
half_ind = np.sort(half_ind)
for index in half_ind:
img_dir = imgs[index]
with open("cmc_half.train", 'a') as f:
f.write(img_dir[len(seq_root):] + '\n')
#
def gen_train_on_dataset(dataset="CMC4"):
seq_root = '../../data/images/'
filename = "cmc_"+dataset+".train"
if os.path.exists(filename):
os.remove(filename)
for cam_num in [1, 2, 3, 4]:
dataset_dir = osp.join(seq_root, dataset, "Cam_" + str(cam_num) + "/*.png")
imgs = glob.glob(dataset_dir)
for img_dir in imgs:
with open(filename, 'a') as f:
f.write(img_dir[len(seq_root):] + '\n')
if __name__ == '__main__':
gen_image_paths()
gen_image_paths_half()
# labels_with_ids()
gen_train_on_dataset()