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()
