modules/softcascade/misc/scale_inria.py

   1 #!/usr/bin/env python
   2
   3 import sys, os, os.path, glob, math, cv2
   4 from datetime import datetime
   5 from optparse import OptionParser
   6
   7 def parse(ipath, f):
   8     bbs = []
   9     path = None
  10     for l in f:
  11         box = None
  12         if l.startswith("Bounding box"):
  13             b = [x.strip() for x in l.split(":")[1].split("-")]
  14             c = [x[1:-1].split(",") for x in b]
  15             d = [int(x) for x in sum(c, [])]
  16             bbs.append(d)
  17
  18         if l.startswith("Image filename"):
  19             path = os.path.join(os.path.join(ipath, ".."), l.split('"')[-2])
  20
  21     return (path, bbs)
  22
  23 def adjust(box, tb, lr):
  24
  25     mix = int(round(box[0] - lr))
  26     miy = int(round(box[1] - tb))
  27
  28     max = int(round(box[2] + lr))
  29     may = int(round(box[3] + tb))
  30
  31     return [mix, miy, max, may]
  32
  33 if __name__ == "__main__":
  34     parser = OptionParser()
  35     parser.add_option("-i", "--input", dest="input", metavar="DIRECTORY", type="string",
  36                        help="path to Inria train data folder")
  37
  38     parser.add_option("-o", "--output", dest="output", metavar="DIRECTORY", type="string",
  39                        help="path to store data", default=".")
  40
  41     parser.add_option("-t", "--target", dest="target", type="string", help="should be train or test", default="train")
  42
  43     (options, args) = parser.parse_args()
  44     if not options.input:
  45         parser.error("Inria data folder required")
  46
  47     if options.target not in ["train", "test"]:
  48         parser.error("dataset should contain train or test data")
  49
  50     octaves = [-1, 0, 1, 2]
  51
  52     path = os.path.join(options.output, datetime.now().strftime("rescaled-" + options.target + "-%Y-%m-%d-%H-%M-%S"))
  53     os.mkdir(path)
  54
  55     neg_path = os.path.join(path, "neg")
  56     os.mkdir(neg_path)
  57
  58     pos_path = os.path.join(path, "pos")
  59     os.mkdir(pos_path)
  60
  61     print "rescaled Inria training data stored into", path, "\nprocessing",
  62     for each in octaves:
  63         octave = 2**each
  64
  65         whole_mod_w = int(64 * octave) + 2 * int(20 * octave)
  66         whole_mod_h = int(128 * octave) + 2 * int(20 * octave)
  67
  68         cpos_path = os.path.join(pos_path, "octave_%d" % each)
  69         os.mkdir(cpos_path)
  70         idx = 0
  71
  72         gl = glob.iglob(os.path.join(options.input, "annotations/*.txt"))
  73         for image, boxes in [parse(options.input, open(__p)) for __p in gl]:
  74             for box in boxes:
  75                 height = box[3] - box[1]
  76                 scale = height / float(96)
  77
  78                 mat = cv2.imread(image)
  79                 mat_h, mat_w, _ = mat.shape
  80
  81                 rel_scale = scale / octave
  82
  83                 d_w = whole_mod_w * rel_scale
  84                 d_h = whole_mod_h * rel_scale
  85
  86                 top_bottom_border = (d_h - (box[3] - box[1])) / 2.0
  87                 left_right_border = (d_w - (box[2] - box[0])) / 2.0
  88
  89                 box = adjust(box, top_bottom_border, left_right_border)
  90                 inner = [max(0, box[0]), max(0, box[1]), min(mat_w, box[2]), min(mat_h, box[3]) ]
  91
  92                 cropped = mat[inner[1]:inner[3], inner[0]:inner[2], :]
  93
  94                 top     = int(max(0, 0 - box[1]))
  95                 bottom  = int(max(0, box[3] - mat_h))
  96                 left    = int(max(0, 0 - box[0]))
  97                 right   = int(max(0, box[2] - mat_w))
  98                 cropped = cv2.copyMakeBorder(cropped, top, bottom, left, right, cv2.BORDER_REPLICATE)
  99                 resized = sft.resize_sample(cropped, whole_mod_w, whole_mod_h)
 100
 101                 out_name = ".png"
 102                 if round(math.log(scale)/math.log(2)) < each:
 103                     out_name = "_upscaled" + out_name
 104
 105                 cv2.imwrite(os.path.join(cpos_path, "sample_%d" % idx + out_name), resized)
 106
 107                 flipped = cv2.flip(resized, 1)
 108                 cv2.imwrite(os.path.join(cpos_path, "sample_%d" % idx + "_mirror" + out_name), flipped)
 109                 idx = idx + 1
 110                 print "." ,
 111                 sys.stdout.flush()
 112
 113         idx = 0
 114         cneg_path = os.path.join(neg_path, "octave_%d" % each)
 115         os.mkdir(cneg_path)
 116
 117         for each in [__n for __n in glob.iglob(os.path.join(options.input, "neg/*.*"))]:
 118             img = cv2.imread(each)
 119             min_shape = (1.5 * whole_mod_h, 1.5 * whole_mod_w)
 120
 121             if (img.shape[1] <= min_shape[1]) or (img.shape[0] <= min_shape[0]):
 122                 out_name = "negative_sample_%i_resized.png" % idx
 123
 124                 ratio = float(img.shape[1]) / img.shape[0]
 125
 126                 if (img.shape[1] <= min_shape[1]):
 127                     resized_size = (int(min_shape[1]), int(min_shape[1] / ratio))
 128
 129                 if (img.shape[0] <= min_shape[0]):
 130                     resized_size = (int(min_shape[0] * ratio), int(min_shape[0]))
 131
 132                 img = sft.resize_sample(img, resized_size[0], resized_size[1])
 133             else:
 134                 out_name = "negative_sample_%i.png" % idx
 135
 136             cv2.imwrite(os.path.join(cneg_path, out_name), img)
 137             idx = idx + 1
 138             print "." ,
 139             sys.stdout.flush()