def saveMasks(seq_path, xml_path, out_mask_size, out_border, fixed_ar, save_raw_mask, show_img, out_root_path='', save_test=1, save_train=1, frames_reader=None, masks_per_seq=0, enable_out_suffix=1, train_fnames=None, test_fnames=None, map_to_bbox=0, out_img_dir='', enable_xml_annotations=0, allow_skipping_images=0): global _pause, _exit if not xml_path or not os.path.isdir(xml_path): raise IOError( 'Folder containing the loaded boxes does not exist: {}'.format( xml_path)) files = glob.glob(os.path.join(xml_path, '*.xml')) n_files = len(files) if n_files == 0: raise IOError('No loaded boxes found') if frames_reader is None: frames_reader = get_frames_reader(seq_path, save_as_bin=False) min_dim = max_dim = 0 out_w, out_h = out_mask_size print('out_mask_size: {}'.format(out_mask_size)) if out_w == -1 and out_h == -1: out_w = out_h = 0 if out_w == -1: max_dim = out_h elif out_h == -1: min_dim = out_w if fixed_ar: print('Using fixed aspect ratio: {}'.format(fixed_ar)) print('out_border: {}'.format(out_border)) def getint(fn): basename = os.path.basename(fn) num = re.sub("\D", "", basename) try: return int(num) except: return 0 if len(files) > 0: files = sorted(files, key=getint) print('Loading annotations from {:d} files'.format(n_files)) file_id = 0 n_boxes = 0 seq_root_dir = os.path.dirname(seq_path) seq_name = os.path.basename(seq_path) if not out_root_path: out_root_path = os.path.join(seq_root_dir, 'masks') if not enable_out_suffix: out_seq_name = seq_name else: if map_to_bbox: out_seq_name = '{}_mapped'.format(seq_name) else: out_seq_name = '{}_{}x{}'.format(seq_name, out_w, out_h) if fixed_ar: out_seq_name = '{}_ar_{}'.format(out_seq_name, fixed_ar) else: out_seq_name = '{}_{}'.format(out_seq_name, out_border) out_seq_name = out_seq_name.replace('.', 'p') train_root_path = os.path.join(out_root_path, out_seq_name) if not save_test and not save_train: raise AssertionError('Either save_test or save_train must be on') # print('Saving output sequences to {}'.format(out_root_path)) if save_train: out_img_root_path = train_root_path if out_img_dir: out_img_root_path = os.path.join(out_img_root_path, out_img_dir) out_mask_root_path = os.path.join(train_root_path, 'labels') print('Saving training mask sequence to {}'.format(train_root_path)) if not os.path.isdir(out_img_root_path): os.makedirs(out_img_root_path) if not os.path.isdir(out_mask_root_path): os.makedirs(out_mask_root_path) if enable_xml_annotations: out_xml_path = os.path.join(out_img_root_path, 'annotations') print('Saving xml_annotations to {}'.format(out_xml_path)) if not os.path.isdir(out_xml_path): os.makedirs(out_xml_path) if save_test: out_test_seq_name = out_seq_name + '_test' test_img_root_path = os.path.join(out_root_path, out_test_seq_name) print('Saving unlabeled testing mask sequence to {}'.format( test_img_root_path)) if not os.path.isdir(test_img_root_path): os.makedirs(test_img_root_path) win_name = 'patch and mask' disable_resizing = 0 scale_x = scale_y = 1.0 if out_w == 0 and out_h == 0: print('Resizing disabled') disable_resizing = 1 csv_raw = [] test_csv_raw = [] n_files = len(files) if save_raw_mask: print('Saving raw labels') mask_pix_val = (1, 1, 1) else: mask_pix_val = (255, 255, 255) n_masks = 0 _train_fnames = [] _test_fnames = [] _exit_seq = 0 disp_img = None for file_id, file in enumerate(files): xml_reader = PascalVocReader(file) filename = os.path.basename(xml_reader.filename) filename_no_ext = os.path.splitext(filename)[0] # file_id = int(re.sub("\D", "", filename)) # print('filename: {}'.format(filename)) # print('file_id: {}'.format(file_id)) img = frames_reader.get_frame_by_name(filename, convert_to_rgb=0) if img is None: print('image {} could not be read'.format(filename)) continue img_h, img_w = img.shape[:2] mask_img = None shapes = xml_reader.getShapes() n_shapes = len(shapes) if n_shapes > 1: print('{} boxes found for {} in {}'.format(n_shapes, filename, file)) obj_id = 0 img_written = 0 for shape in shapes: label, points, _, _, difficult, bbox_source, id_number, score, mask, mask_img = shape if not mask: if not save_test: continue xmin, ymin = points[0] xmax, ymax = points[2] img_root_path = test_img_root_path else: if not save_train: continue mask_pts_list = Shape.getContourPts(mask, verbose=0) mask_pts = np.asarray(mask_pts_list) xmin, ymin = np.min(mask_pts, axis=0).astype(np.int32) xmax, ymax = np.max(mask_pts, axis=0).astype(np.int32) img_root_path = out_img_root_path if fixed_ar: w, h = xmax - xmin, ymax - ymin src_ar = float(w) / float(h) if fixed_ar > src_ar: border_x = int((h * fixed_ar - w) / 2.0) border_y = 0 else: border_y = int((w / fixed_ar - h) / 2.0) border_x = 0 else: border_x = border_y = out_border # start_row, start_col = max(0, ymin - border_y), max(0, xmin - border_x) # end_row, end_col = min(img_h - 1, ymax + border_y), min(img_w - 1, xmax + border_x) start_row, start_col = ymin - border_y, xmin - border_x end_row, end_col = ymax + border_y, xmax + border_x if start_row < 0 or start_col < 0 or end_row >= img_h or end_col >= img_w: msg = 'Invalid border {} for box {} in image {} of size {}'.format( [border_x, border_y], [xmin, ymin, xmax, ymax], filename, [img_w, img_h]) if allow_skipping_images: print('\n' + msg + '\n') continue else: raise AssertionError(msg) if mask: n_masks += 1 w, h = end_col - start_col, end_row - start_row patch_img = img[start_row:end_row, start_col:end_col, :] if not disable_resizing: if max_dim > 0: if w > h: out_w = max_dim out_h = 0 else: out_h = max_dim out_w = 0 elif min_dim > 0: if w < h: out_w = min_dim out_h = 0 else: out_h = min_dim out_w = 0 else: out_w, out_h = out_mask_size scale_x = float(out_w) / float(w) scale_y = float(out_h) / float(h) if scale_x == 0: scale_x = scale_y out_w = int(w * scale_x) elif scale_y == 0: scale_y = scale_x out_h = int(h * scale_y) try: patch_img = cv2.resize(patch_img, (out_w, out_h)) # print('patch_img: {}'.format(patch_img.shape)) except cv2.error as e: print('patch_img: {}'.format(patch_img.shape)) print('out_size: {}, {}'.format(start_row, start_col)) print('out_size: {}, {}'.format(end_row, end_col)) print('out_size: {}, {}'.format(out_w, out_h)) raise cv2.error(e) else: out_w, out_h = w, h _label = label if id_number is None: id_number = -1 if id_number > 0: _label = '{}_{}'.format(_label, id_number) if enable_out_suffix: out_fname = '{}_{}_{}'.format(filename_no_ext, obj_id, label) else: out_fname = filename_no_ext _xmin, _ymin = int((xmin - start_col) * scale_x), int( (ymin - start_row) * scale_y) _xmax, _ymax = int((xmax - start_col) * scale_x), int( (ymax - start_row) * scale_y) if map_to_bbox: if not img_written: img_written = 1 out_img_path = os.path.join(img_root_path, filename) cv2.imwrite(out_img_path, img) if enable_xml_annotations: imageShape = [xml_reader.height, xml_reader.width, 3] xml_writer = PascalVocWriter(out_xml_path, filename, imageShape) if mask: if enable_xml_annotations: bndbox = [xmin, ymin, xmax, ymax] xml_writer.addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label, difficult, bbox_source, id_number, score, mask, mask_img) raw_data = { 'target_id': int(id_number), 'filename': filename, 'width': img_w, 'height': img_h, 'class': label, 'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax } if show_img: cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2) disp_img = img else: img_out_fname = out_fname + '.jpg' if mask: out_img_path = os.path.join(img_root_path, img_out_fname) cv2.imwrite(out_img_path, patch_img) if enable_xml_annotations: n_mask = len(mask) _mask = [] for i in range(n_mask): _mask.append([(mask[i][0] - start_col) * scale_x, (mask[i][1] - start_row) * scale_y, mask[i][2]]) imageShape = [xml_reader.height, xml_reader.width, 3] xml_writer = PascalVocWriter(out_xml_path, xml_reader.filename, imageShape) bndbox = [_xmin, _ymin, _xmax, _ymax] xml_writer.addBndBox(_xmin, _ymin, _xmax, _ymax, label, difficult, bbox_source, id_number, score, _mask) raw_data = { 'target_id': int(id_number), 'filename': img_out_fname, 'width': out_w, 'height': out_h, 'class': label, 'xmin': _xmin, 'ymin': _ymin, 'xmax': _xmax, 'ymax': _ymax } if show_img: cv2.rectangle(patch_img, (_xmin, _ymin), (_xmax, _ymax), (0, 255, 0), 2) disp_img = patch_img if mask: if mask_img is None: mask_img = np.zeros_like(img) # print('border_x: {}'.format(border_x)) # print('border_y: {}'.format(border_y)) # print('scale_x: {}'.format(scale_x)) # print('scale_y: {}'.format(scale_y)) # # print('xmin: {}'.format(xmin)) # print('ymin: {}'.format(ymin)) mask_pts = [[(x - xmin + border_x) * scale_x, (y - ymin + border_y) * scale_y] for x, y in mask_pts] curr_mask = np.zeros_like(patch_img, dtype=np.uint8) # print('mask_img: {}'.format(mask_img.shape)) mask_out_fname = out_fname + '.png' # np.savetxt('mask_seq_mask_pts.txt', mask_pts, fmt='%.6f') curr_mask = cv2.fillPoly( curr_mask, np.array([ mask_pts, ], dtype=np.int32), mask_pix_val) # print('min: {} max: {}'.format( # np.min(mask_img.flatten()), # np.max(mask_img.flatten())) # ) if map_to_bbox: mask_img = map_mask_to_bbox( (xmin, ymin, xmax, ymax), curr_mask, fixed_ar, out_border, mask_img.shape, mask_img) else: mask_img = curr_mask out_mask_path = os.path.join(out_mask_root_path, mask_out_fname) cv2.imwrite(out_mask_path, mask_img) _train_fnames.append((out_img_path, out_mask_path)) if show_img: disp_mask_img = mask_img.copy() if save_raw_mask: disp_mask_img[disp_mask_img > 0] = 255 blended_img = np.asarray( Image.blend(Image.fromarray(patch_img), Image.fromarray(disp_mask_img), 0.5)) disp_img = np.concatenate( (disp_img, disp_mask_img, blended_img), axis=1) csv_raw.append(raw_data) else: test_csv_raw.append(raw_data) if not map_to_bbox: _test_fnames.append(out_img_path) if show_img and not map_to_bbox: # if _pause: # print('frame {} :: {}'.format(file_id, filename)) cv2.imshow(win_name, disp_img) k = cv2.waitKey(1 - _pause) if k == ord('q'): _exit = 1 break elif k == 27: _exit_seq = 1 break elif k == 32: _pause = 1 - _pause obj_id += 1 if map_to_bbox and img is not None: out_img_path = os.path.join(out_img_root_path, filename) if save_train and mask_img is not None: mask_out_fname = filename_no_ext + '.png' out_mask_path = os.path.join(out_mask_root_path, mask_out_fname) cv2.imwrite(out_mask_path, mask_img) if enable_xml_annotations: out_xml_file = os.path.join(out_xml_path, os.path.basename(file)) xml_writer.save(targetFile=out_xml_file) _train_fnames.append((out_img_path, out_mask_path)) if show_img: disp_mask_img = mask_img if save_raw_mask: disp_mask_img[disp_mask_img > 0] = 255 blended_img = np.asarray( Image.blend(Image.fromarray(img), Image.fromarray(disp_mask_img), 0.5)) disp_img = np.concatenate( (disp_img, disp_mask_img, blended_img), axis=1) elif save_test: out_img_path = os.path.join(test_img_root_path, filename) if out_img_path in _test_fnames: raise IOError( 'Duplicate out_img_path: {}'.format(out_img_path)) _test_fnames.append(out_img_path) if show_img and disp_img is not None: cv2.imshow(win_name, disp_img) k = cv2.waitKey(1 - _pause) if k == ord('q'): _exit = 1 break elif k == 27: break elif k == 32: _pause = 1 - _pause if _exit: break sys.stdout.write( '\rDone {:d}/{:d} files {:s} ({:d} masks found)'.format( file_id + 1, n_files, filename, n_masks)) sys.stdout.flush() if masks_per_seq > 0 and n_masks >= masks_per_seq: break sys.stdout.write('\n') sys.stdout.flush() if not _exit_seq and save_train and n_masks == 0: raise IOError('\nNo masks found for {}\n'.format(seq_path)) train_csv_path = test_csv_path = '' if csv_raw: print('Saved {} labeled files in training sequence'.format( len(csv_raw))) train_csv_path = os.path.join(out_img_root_path, 'annotations.csv') pd.DataFrame(csv_raw).to_csv(train_csv_path) if test_csv_raw: print('Saved {} unlabeled files in test sequence'.format( len(test_csv_raw))) test_csv_path = os.path.join(test_img_root_path, 'annotations.csv') pd.DataFrame(test_csv_raw).to_csv(test_csv_path) if show_img: cv2.destroyWindow(win_name) if save_train and train_fnames is not None: train_fnames[ out_seq_name] = _train_fnames, train_root_path, csv_raw, train_csv_path if save_test and test_fnames is not None: test_fnames[ out_test_seq_name] = _test_fnames, test_img_root_path, test_csv_raw, test_csv_path return n_masks
def visualize(vis_params, logger, img_path, csv_path, class_dict, init_frame_id=0, n_frames=0, request_roi=None, generator_mode=0, enable_masks=0, label='', only_boxes=0, crop_size=()): """ :param vis_params: :param logger: :param img_path: :param csv_path: :param class_dict: :param init_frame_id: :param n_frames: :param request_roi: :param generator_mode: :return: """ global _pause, _quit save_fname_templ = os.path.splitext(os.path.basename(img_path))[0] # csv_path = os.path.join(img_path, 'annotations.csv') df = pd.read_csv(csv_path) frames_reader = get_frames_reader(img_path, save_as_bin=False) if request_roi is not None: frames_reader.setROI(request_roi) class_labels = dict((v, k) for k, v in class_dict.items()) if generator_mode: vis_params.show = 0 vis_params.save = 0 visualizer = Visualizer(vis_params, logger, class_labels) init_frame = frames_reader.get_frame(init_frame_id) height, width, _ = init_frame.shape frame_size = width, height visualizer.initialize(save_fname_templ, frame_size, _pause) if n_frames <= 0: n_frames = frames_reader.num_frames print('Reading {:d} images from {:s}...'.format(n_frames, img_path)) for frame_id in range(init_frame_id, n_frames): try: curr_frame = frames_reader.get_frame(frame_id) except IOError as e: print('{}'.format(e)) break if only_boxes: curr_frame = np.zeros_like(curr_frame) file_path = frames_reader.get_file_path() if file_path is None: print('Visualization is only supported on image sequence data') return filename = os.path.basename(file_path) multiple_instance = df.loc[df['filename'] == filename] # Total # of object instances in a file n_bboxes = len(multiple_instance.index) # Remove from df (avoids duplication) df = df.drop(multiple_instance.index[:n_bboxes]) frame_data = [] masks = [] generic_target_id = -1 if enable_masks: filename = os.path.basename(file_path) xml_path = os.path.join(img_path, 'annotations', os.path.splitext(filename)[0] + '.xml') if not os.path.isfile(xml_path): print('{} :: annotations xml file not found: {}'.format( filename, xml_path)) continue xml_reader = PascalVocReader(xml_path) shapes = xml_reader.getShapes() n_shapes = len(shapes) if n_shapes != n_bboxes: raise IOError( 'Mismatch between n_bboxes in xml: {} and csv: {}'.format( n_shapes, n_bboxes)) for box_id in range(n_bboxes): bbox = multiple_instance.iloc[box_id] try: target_id = bbox['target_id'] except KeyError: target_id = generic_target_id generic_target_id -= 1 xmin = bbox.loc['xmin'] ymin = bbox.loc['ymin'] xmax = bbox.loc['xmax'] ymax = bbox.loc['ymax'] class_name = bbox.loc['class'] try: class_id = class_dict[str(class_name)] except KeyError: print('Ignoring annotation with invalid class: {}'.format( class_name)) continue width = xmax - xmin height = ymax - ymin curr_frame_data = [ frame_id, target_id, xmin, ymin, width, height, class_id ] if enable_masks: mask = shapes[box_id][-2] if mask is not None: _contour_pts = Shape.getContourPts(mask) masks.append(_contour_pts) frame_data.append(curr_frame_data) frame_data = np.asarray(frame_data) res = visualizer.update(frame_id, curr_frame, frame_data, masks, label, crop_size) if generator_mode: yield res # elif not res: # break _quit = visualizer._quit _pause = visualizer._pause visualizer.close() frames_reader.close()