Crowd-Counting-and-Localization/dataset_preparation.py at main · HoseinRanjbar/Crowd-Counting-and-Localization · GitHub

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
import argparse
import numpy as np
import cv2
import os
from scipy.io import loadmat

def get_args_parser():
    parser = argparse.ArgumentParser('dataset preparation', add_help=False)

    parser.add_argument('--application', default='both',
                        help='make_list , convert_gr_format or both')

    parser.add_argument('--groundtruth_dir', default='',
                        help='path where to load groundtruth')

    parser.add_argument('--groundtruth_txt_path', default='',
                        help='path where to save groundtruth information with .txt format')

    parser.add_argument('--list_dir', default='',
                        help='path where to save')

    parser.add_argument('--images_path', default='')

    return parser

def main(args):
  gt_path=args.groundtruth_dir
  gt_txt_path=args.groundtruth_txt_path
  list_path=args.list_dir
  imgs_path=args.images_path
  app=args.application

  if app=='convert_gr_format':
    #convert grandtruth information from .mat to .txt
    for filename in os.listdir(gt_path):
      gt = loadmat(os.path.join(gt_path,filename))
      info=(gt['image_info'])[0][0][0][0][0]
      file = open(os.path.join(gt_txt_path,'{}.txt'.format(filename[:-4])), "w+")
      for i in range(np.shape(info)[0]):
        x,y=info[i][0],info[i][1]
        file.write("{} {}\n".format(x,y))

  elif app=='make_list':
    #Creat list of images and grandtruth for training
    file = open(os.path.join(list_path,'train_list.txt'), "w+")
    for filename in os.listdir(imgs_path):
      file.write("images/{} gt_txt/GT_{}.txt\n".format(filename,filename[:-4]))

  else:
    #convert grandtruth information from .mat to .txt
    for filename in os.listdir(gt_path):
      gt = loadmat(os.path.join(gt_path,filename))
      info=(gt['image_info'])[0][0][0][0][0]
      file = open(os.path.join(gt_txt_path,'{}.txt'.format(filename[:-4])), "w+")
      for i in range(np.shape(info)[0]):
        x,y=info[i][0],info[i][1]
        file.write("{} {}\n".format(x,y))

    #Creat list of images and grandtruth for training
    file = open(os.path.join(list_path,'train_list.txt'), "w+")
    for filename in os.listdir(imgs_path):
      file.write("images/{} gt_txt/GT_{}.txt\n".format(filename,filename[:-4]))


if __name__ == '__main__':
    parser = argparse.ArgumentParser('dataset preparation', parents=[get_args_parser()])
    args = parser.parse_args()
    main(args)