def download_images()

in scripts/download_data.py [0:0]

• 94 lines of code
• 33 McCabe index (conditional complexity)

  def download_images(self, p_num, intent, dload_dir,
                      include_objects=None, download_color=True,
                      download_depth=True):
    assert osp.isdir(dload_dir),\
      'Image download dir {:s} does not exist'.format(dload_dir)
    p_id = 'full{:d}_{:s}'.format(p_num, intent)
    if download_color and (not download_depth):
      urls = self.urls['videos']['color']
    else:
      urls = self.urls['images']
    
    # check if already extracted
    dirs_to_check = []
    if download_color:
      dirs_to_check.append('color')
    if download_depth:
      dirs_to_check.append('depth')
    ok = True
    if osp.isdir(osp.join(self.data_dir, p_id)):
      sess_dir = osp.join(self.data_dir, p_id)
      for object_name in next(os.walk(sess_dir))[1]:
        if include_objects is not None and object_name not in include_objects:
          continue
        images_dir = osp.join(sess_dir, object_name, 'images_full')
        if not osp.isdir(images_dir):
          continue
        for cam_name in next(os.walk(images_dir))[1]:
          for check_name in dirs_to_check:
            check_dir = osp.join(images_dir, cam_name, check_name)
            if is_nonempty_dir(check_dir):
              print('{:s} {:s} already has extracted images, please delete {:s}'.
                    format(p_id, object_name, check_dir))
              ok = False
    if not ok:
      return
    
    # download and extract
    sess_dir = osp.join(dload_dir, p_id)
    if not osp.isdir(sess_dir):
      print('Creating {:s}'.format(sess_dir))
    os.makedirs(sess_dir, exist_ok=True)
    print('Downloading {:s} images...'.format(p_id))
    object_names = list(urls[p_id].keys())
    if include_objects is None:
      include_objects = object_names[:]
    filenames_to_extract = {}
    for object_name in tqdm(include_objects):
      if object_name not in object_names:
        print('{:d} {:s} does not have {:s}'.format(p_num, intent, object_name))
        continue
      filename = osp.join(sess_dir, '{:s}_images.zip'.format(object_name))
      url = urls[p_id][object_name]
      print(object_name)
      if nutils.download_url(url, filename):
        filenames_to_extract[object_name] = filename
      else:
        print('{:s} {:s} Download unsuccessful'.format(p_id, object_name))
        return
    
    print('Extracting...')
    for object_name, filename in tqdm(filenames_to_extract.items()):
      obj_dir = osp.join(sess_dir, object_name)
      os.makedirs(obj_dir, exist_ok=True)
      self._unzip_and_del(filename, obj_dir)
      for filename in next(os.walk(obj_dir))[-1]:
        if download_color and (not download_depth):
          if '.mp4' not in filename:
            continue
          camera_name = filename.replace('.mp4', '')
          video_filename = osp.join(obj_dir, filename)
          im_dir = osp.join(obj_dir, 'images_full', camera_name, 'color')
          os.makedirs(im_dir, exist_ok=True)
          cap = cv2.VideoCapture(video_filename)
          if not cap.isOpened():
            print('Could not read {:s}'.format(video_filename))
            return
          count = 0
          while True:
            ok, im = cap.read()
            if not ok:
              break
            filename = osp.join(im_dir, 'frame{:03d}.png'.format(count))
            cv2.imwrite(filename, im)
            count += 1
          os.remove(video_filename)
        else:
          if '.zip' not in filename:
            continue
          filter_fn = (lambda x: 'color' not in x) if (not download_color) \
              else None
          self._unzip_and_del(osp.join(obj_dir, filename), progress=False,
                              filter_fn=filter_fn)

      # symlink
      if osp.realpath(dload_dir) != osp.realpath(self.data_dir):
        src = osp.join(obj_dir, 'images_full')
        dst_dir = osp.join(self.data_dir, p_id, object_name)
        if not osp.isdir(dst_dir):
          os.makedirs(dst_dir)
        dst = osp.join(dst_dir, 'images_full')
        os.symlink(src, dst)