Meta Byte Track
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

generate_half_csv.py 1.5KB

12345678910111213141516171819202122232425262728293031323334353637
  1. import os
  2. import numpy as np
  3. prefix_dir = 'MOT17/'
  4. root_dir = 'train/'
  5. result_csv = 'train_half_annots.csv'
  6. train_half_set = {2: 301, 4: 526, 5:419, 9:263, 10:328, 11:451, 13:376}
  7. fout = open(result_csv, 'w')
  8. for data_name in sorted(os.listdir(prefix_dir + root_dir)):
  9. print(data_name)
  10. gt_path = os.path.join(prefix_dir, root_dir, data_name, 'gt', 'gt.txt')
  11. # print(gt_path)
  12. data_raw = np.loadtxt(gt_path, delimiter=',', dtype='float', usecols=(0,1,2,3,4,5,6,7,8))
  13. data_sort = data_raw[np.lexsort(data_raw[:,::-1].T)]
  14. visible_raw = data_sort[:,8]
  15. # print(data_sort)
  16. # print(data_sort[-1, 0])
  17. img_num = data_sort[-1, 0]
  18. # print(data_sort.shape[0])
  19. box_num = data_sort.shape[0]
  20. person_box_num = np.sum(data_sort[:,6] == 1)
  21. # print(person_box_num)
  22. # import ipdb; ipdb.set_trace()
  23. for i in range(box_num):
  24. c = int(data_sort[i, 6])
  25. v = visible_raw[i]
  26. img_index = int(data_sort[i, 0])
  27. if c == 1 and v > 0.1 and img_index < train_half_set[int(data_name[-2:])]:
  28. img_index = int(data_sort[i, 0])
  29. img_name = data_name + '/img1/' + str(img_index).zfill(6) + '.jpg'
  30. print(root_dir + img_name + ', ' + str(int(data_sort[i, 1])) + ', ' + str(data_sort[i, 2]) + ', ' + str(data_sort[i, 3]) + ', ' + str(data_sort[i, 2] + data_sort[i, 4]) + ', ' + str(data_sort[i, 3] + data_sort[i, 5]) + ', person\n')
  31. fout.write(root_dir + img_name + ', ' + str(int(data_sort[i, 1])) + ', ' + str(data_sort[i, 2]) + ', ' + str(data_sort[i, 3]) + ', ' + str(data_sort[i, 2] + data_sort[i, 4]) + ', ' + str(data_sort[i, 3] + data_sort[i, 5]) + ', person\n')
  32. fout.close()