def fixMasks(img_path, mask_path, csv_path, out_size='', out_dir='',
save_video=1, write_text=0, show_img=1, frames_reader=None,
img_ext='jpg', mask_ext='png', out_ext='mkv', codec='H264', fps=30, border=0, fixed_ar=0,
include_binary=1, include_orig=1, show_bbox=1, map_to_bbox=0, apply_contour=0,
writer=None, raw_mask=0):
global _pause, _exit
if map_to_bbox and not csv_path:
csv_path = os.path.join(img_path, 'annotations.csv')
img_files = [os.path.join(img_path, k) for k in os.listdir(img_path) if
os.path.splitext(k.lower())[1][1:] == img_ext]
mask_files = [os.path.join(mask_path, k) for k in os.listdir(mask_path) if
os.path.splitext(k.lower())[1][1:] == mask_ext]
n_img_files, n_mask_files = len(img_files), len(mask_files)
if n_img_files == 0:
raise IOError('No img_files of type {} found in {}'.format(img_ext, img_path))
if n_mask_files == 0:
raise IOError('No mask_files of type {} found in {}'.format(mask_ext, mask_path))
print('Found {} image and {} mask files'.format(n_img_files, n_mask_files))
win_name = 'patch and mask'
text_fmt = ('green', 0, 5, 1.0, 1)
text_color = utils.col_rgb[text_fmt[0]]
text_font = utils.CVConstants.fonts[text_fmt[2]]
text_font_size = text_fmt[3]
text_thickness = text_fmt[4]
text_location = (5, 15)
if cv2.__version__.startswith('2'):
font_line_type = cv2.CV_AA
else:
font_line_type = cv2.LINE_AA
# n_disp_img = 1
# if include_orig:
# n_disp_img += 1
# if include_binary:
# n_disp_img += 1
retrieval_mode = None
if apply_contour == 1:
retrieval_mode = cv2.RETR_EXTERNAL
elif apply_contour == 2:
retrieval_mode = cv2.RETR_CCOMP
elif apply_contour == 3:
retrieval_mode = cv2.RETR_TREE
out_path = ''
if save_video:
out_name = os.path.basename(mask_path)
if apply_contour:
out_name += '_contour_{}'.format(apply_contour)
out_name += '.' + out_ext
if not out_dir:
out_dir = os.path.dirname(mask_path)
out_path = os.path.join(out_dir, out_name)
seq_name = os.path.basename(img_path)
print('img_path: ', img_path)
print('mask_path: ', mask_path)
print('csv_path: ', csv_path)
print('seq_name: ', seq_name)
# if n_img_files != n_mask_files:
# raise IOError('Mismatch between n_img_files: {} and n_mask_files: {}'.format(n_img_files, n_mask_files))
one_to_one_mode = 0
if n_img_files == n_mask_files:
print('Using one_to_one_mode instead of filename matching to associate masks with images')
one_to_one_mode = 1
if csv_path:
df = pd.read_csv(csv_path)
if out_path and writer is None:
if out_size:
out_size = tuple([int(x) for x in out_size.split('x')])
if 0 in out_size:
img_h, img_w = cv2.imread(img_files[0]).shape[:2]
if out_size[0] == 0 and out_size[1] == 0:
out_w, out_h = img_w, img_h
elif out_size[0] == 0:
out_h = out_size[1]
out_w = int(img_w * (float(out_h) / float(img_h)))
elif out_size[1] == 0:
out_w = out_size[0]
out_h = int(img_h * (float(out_w) / float(img_w)))
out_size = (out_w, out_h)
else:
out_size = (1920, 1080)
image_exts = ['jpg', 'bmp', 'png']
if out_ext[0] in image_exts:
writer = ImageWriter(out_path)
else:
writer = cv2.VideoWriter()
writer_params = {
'filename': out_path,
'fps': int(fps),
'frameSize': out_size,
}
if cv2.__version__.startswith('2'):
writer_params['fourcc'] = cv2.cv.CV_FOURCC(*codec)
else:
writer_params['apiPreference'] = cv2.CAP_FFMPEG
writer_params['fourcc'] = cv2.VideoWriter_fourcc(*codec)
writer.open(**writer_params)
if not writer.isOpened():
raise IOError('Video file {:s} could not be opened'.format(out_path))
print('Writing {}x{} video to {}'.format(out_size[0], out_size[1], out_path))
n_valid_masks = 0
for file_id in range(n_img_files):
img_file = img_files[file_id]
img_file = img_file.replace('\\', '/')
filename = os.path.basename(img_file)
filename_no_ext = os.path.splitext(filename)[0]
if one_to_one_mode:
curr_mask_files = [mask_files[file_id], ]
else:
curr_mask_files = [k for k in mask_files if
os.path.basename(k).startswith(filename_no_ext)]
# filename_no_ext_rev = filename_no_ext[::-1]
# curr_mask_files = [k for k in mask_files if
# os.path.splitext(os.path.basename(k))[0][::-1].startswith(filename_no_ext_rev)
# and os.path.commonprefix([os.path.basename(k), filename_no_ext])]
if not curr_mask_files:
# print('No matching curr_mask file found for {} with filename_no_ext: {}'.format(
# img_file, filename_no_ext))
continue
img = cv2.imread(img_file)
if img is None:
raise IOError('img_file could not be read: {}'.format(img_file))
img_h, img_w = img.shape[:2]
if csv_path:
bboxes = df.loc[df['filename'] == filename]
n_bboxes = len(bboxes.index)
df = df.drop(bboxes.index[:n_bboxes])
if one_to_one_mode and n_bboxes > 1:
print('Considering only one out of {} boxes for {}'.format(n_bboxes, filename))
n_bboxes = 1
n_masks = len(curr_mask_files)
if n_masks != n_bboxes:
raise IOError('Mismatch between n_bboxes: {} and n_masks: {} for {}'.format(
n_bboxes, n_masks, img_file))
mask_img = np.zeros_like(img)
for box_id in range(n_bboxes):
mask_file = curr_mask_files[box_id].replace('\\', '/')
curr_mask = cv2.imread(mask_file)
if curr_mask is None:
raise IOError('mask_file could not be read: {}'.format(mask_file))
if apply_contour:
if apply_contour > 1:
contours, _ = cv2.findContours(curr_mask[:, :, 0].astype(np.uint8), retrieval_mode,
cv2.CHAIN_APPROX_NONE)[-2:]
curr_mask = np.zeros_like(curr_mask, dtype=np.uint8)
cv2.drawContours(curr_mask, contours, -1, (255, 255, 255), -1)
else:
contours, _ = Shape.contourPtsFromMask(curr_mask)
curr_mask, _ = Shape.contourPtsToMask(contours, curr_mask)
mask_h, mask_w = curr_mask.shape[:2]
bbox = bboxes.iloc[box_id]
xmin = bbox.loc['xmin']
ymin = bbox.loc['ymin']
xmax = bbox.loc['xmax']
ymax = bbox.loc['ymax']
mask_img = map_mask_to_bbox((xmin, ymin, xmax, ymax), curr_mask,
fixed_ar, border, mask_img.shape, mask_img)
if show_bbox:
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
if map_to_bbox == 2:
img = img[ymin:ymax, xmin:xmax, :]
mask_img = mask_img[ymin:ymax, xmin:xmax, :]
else:
mask_file = curr_mask_files[0].replace('\\', '/')
mask_img = cv2.imread(mask_file)
if mask_img is None:
raise IOError('mask_file could not be read: {}'.format(mask_file))
mask_h, mask_w = mask_img.shape[:2]
if img_h != mask_h or img_w != mask_w:
mask_img = cv2.resize(mask_img, (img_w, img_h))
if raw_mask:
mask_img *= 255
blended_img = np.asarray(Image.blend(Image.fromarray(img), Image.fromarray(mask_img), 0.5))
def paintMouseHandler(event, x, y, flags=None, param=None):
nonlocal mask_pts, contour_pts, blended_img, mask_img, disp_img, del_thresh, mag_patch_size, mag_win_size, \
mouse_x, mouse_y, prev_mouse_pt, draw_mask_kb, paint_mode
paint_mode = 1
mouse_x, mouse_y = x, y
# refresh_paint_win = 1
draw_marker = 1
marker_col = (0, 255, 0)
# _show_magnified_window = 0
if event == cv2.EVENT_MOUSEMOVE:
# print('flags: {}'.format(flags))
# refresh_paint_win = 1
if flags == 1 or flags == 25:
# left button
min_x, min_y = x - del_thresh, y - del_thresh
max_x, max_y = x + del_thresh, y + del_thresh
mask_img[min_y:max_y, min_x:max_x, :] = 255
blended_img[min_y:max_y, min_x:max_x, :] = (shape_patch[min_y:max_y, min_x:max_x,
:] + 255.0) / 2.0
marker_col = (0, 255, 0)
elif flags == 17:
# shift + left button
min_x, min_y = x - del_thresh, y - del_thresh
max_x, max_y = x + del_thresh, y + del_thresh
mask_img[min_y:max_y, min_x:max_x, :] = 0
blended_img[min_y:max_y, min_x:max_x, :] = (shape_patch[min_y:max_y, min_x:max_x, :]) / 2.0
marker_col = (0, 0, 255)
elif flags == 32:
draw_marker = 0
elif flags == 16:
marker_col = (0, 0, 255)
elif flags == 8:
draw_marker = 0
# else:
# draw_marker = 0
elif event == cv2.EVENT_LBUTTONDOWN:
# print('flags: {}'.format(flags))
pass
elif event == cv2.EVENT_RBUTTONUP:
# print('flags: {}'.format(flags))
pass
elif event == cv2.EVENT_MBUTTONDOWN:
contour_pts, mask_pts = self.contourPtsFromMask(mask_img)
draw_mask_kb = 1
# print('flags: {}'.format(flags))
elif event == cv2.EVENT_MOUSEWHEEL:
if flags > 0:
if flags == mouse_whl_keys_to_flags['ctrl+alt+shift'][0]:
mag_win_size += 10
print('magnified window size increased to {}'.format(mag_win_size))
elif flags == mouse_whl_keys_to_flags['ctrl+shift'][0]:
mag_patch_size -= 1
if mag_patch_size < 5:
mag_patch_size = 5
else:
if del_thresh < 10:
del_thresh += 1
else:
del_thresh += 5
print('del_thresh increased to {}'.format(del_thresh))
else:
if flags == mouse_whl_keys_to_flags['ctrl+alt+shift'][1]:
mag_win_size -= 10
if mag_win_size < 100:
mag_win_size = 100
print('magnified window size decreased to {}'.format(mag_win_size))
elif flags == mouse_whl_keys_to_flags['ctrl+shift'][1]:
mag_patch_size += 1
else:
if del_thresh < 10:
del_thresh = max(del_thresh - 1, 1)
else:
del_thresh -= 5
print('del_thresh decreased to {}'.format(del_thresh))
if draw_marker:
disp_img = np.copy(blended_img)
min_x, min_y = x - del_thresh, y - del_thresh
max_x, max_y = x + del_thresh, y + del_thresh
cv2.rectangle(disp_img, (min_x, min_y), (max_x, max_y), marker_col, 1)
else:
disp_img = blended_img
cv2.imshow(paint_win_name, disp_img)
if show_magnified_window:
showMagnifiedWindow(x, y, _shape_patch, draw_marker=2,
marker_col=marker_col,
# win_name='Paint Magnified'
)
# cv2.imshow('binary mask', mask_img)
prev_mouse_pt = (x, y)
if include_orig and include_binary and 3 * img_w > out_size[0]:
disp_img = np.concatenate((img, mask_img), axis=0)
disp_img = np.concatenate((disp_img, cv2.resize(blended_img, (0, 0), fx=2, fy=2)), axis=1)
else:
disp_img = blended_img
if include_binary:
disp_img = np.concatenate((mask_img, disp_img), axis=1)
if include_orig:
disp_img = np.concatenate((img, disp_img), axis=1)
disp_img = resizeAR(disp_img, out_size[0], out_size[1])
if save_video:
if write_text:
cv2.putText(disp_img, '{} frame {:d}'.format(seq_name, file_id + 1), text_location,
text_font, text_font_size, text_color, text_thickness, text_line_type)
writer.write(disp_img)
if show_img:
cv2.imshow(win_name, disp_img)
k = cv2.waitKey(1 - _pause)
if k == 27:
break
elif k == ord('q'):
_exit = 1
break
elif k == 32:
_pause = 1 - _pause
if _exit:
break
sys.stdout.write('\rDone {:d}/{:d} files'.format(
file_id, n_img_files))
sys.stdout.flush()
n_valid_masks += 1
sys.stdout.write('\n')
sys.stdout.flush()
if n_valid_masks == 0:
raise IOError('No valid masks found')
# if save_video:
# writer.release()
if show_img:
cv2.destroyWindow(win_name)
return writer, out_size
def main():
global _pause, _quit
params = VisParams()
paramparse.process(params)
# _args = [k for k in sys.argv[1:] if not k.startswith('vis.')]
# vis_args = ['--{}'.format(k.replace('vis.', '')) for k in sys.argv[1:] if k.startswith('vis.')]
# processArguments(_args, params)
# params = _params
seq_paths = params.seq_paths
root_dir = params.root_dir
csv_paths = params.csv_paths
csv_root_dir = params.csv_root_dir
class_names_path = params.class_names_path
data_type = params.data_type
n_frames = params.n_frames
seq_prefix = params.seq_prefix
n_vis = params.n_vis
vis_size = params.vis_size
enable_masks = params.enable_masks
show_img = params.show_img
save = params.save
save_fmt = params.save_fmt
save_dir = params.save_dir
labels = params.labels
grid_size = params.grid_size
only_boxes = params.only_boxes
crop_size = params.crop_size
if crop_size:
crop_size = tuple([int(x) for x in crop_size.split('x')])
print('Cropping a region of size {}x{} around the box'.format(
*crop_size))
else:
crop_size = ()
if grid_size:
grid_size = [int(k) for k in grid_size.split('x')]
print('Using a grid size of {}x{}'.format(*grid_size))
else:
grid_size = None
# params = Namespace(**params)
if vis_size:
vis_size = [int(x) for x in vis_size.split('x')]
# get parameters
# _params = ServerParams()
# _params.processArguments()
# print('vis_args: ', _params.vis.__dict__)
# processArguments2(_params, vis_args)
# print('_params: ', _params)
# setup logger
logging_fmt = '%(levelname)s::%(module)s::%(funcName)s::%(lineno)s : %(message)s'
logging_level = logging.INFO
# logging_level = logging.DEBUG
# logging_level = PROFILE_LEVEL_NUM
logging.basicConfig(level=logging_level, format=logging_fmt)
_logger = logging.getLogger()
_logger.setLevel(logging.INFO)
if seq_paths:
if os.path.isfile(seq_paths):
seq_paths = [
x.strip() for x in open(seq_paths).readlines() if x.strip()
]
else:
seq_paths = seq_paths.split(',')
if root_dir:
seq_paths = [os.path.join(root_dir, name) for name in seq_paths]
elif root_dir:
seq_paths = [
os.path.join(root_dir, name) for name in os.listdir(root_dir)
if os.path.isdir(os.path.join(root_dir, name))
]
seq_paths.sort(key=sortKey)
else:
raise IOError('Either seq_paths or root_dir must be provided')
if csv_paths:
if os.path.isfile(csv_paths):
csv_paths = [
x.strip() for x in open(csv_paths).readlines() if x.strip()
]
else:
csv_paths = csv_paths.split(',')
if csv_root_dir:
csv_paths = [
os.path.join(csv_root_dir, name) for name in csv_paths
]
elif csv_root_dir:
csv_paths = [
os.path.join(csv_root_dir, name)
for name in os.listdir(csv_root_dir)
if os.path.isfile(os.path.join(csv_root_dir, name))
and name.endswith('.csv')
]
csv_paths.sort(key=sortKey)
else:
csv_paths = [
os.path.join(seq_path, data_type + '.csv')
for seq_path in seq_paths
]
seq_path_ids = []
if seq_prefix:
seq_path_ids = [
_id for _id, seq_path in enumerate(seq_paths)
if os.path.basename(seq_path).startswith(seq_prefix)
]
seq_paths = [seq_paths[_id] for _id in seq_path_ids]
csv_paths = [csv_paths[_id] for _id in seq_path_ids]
n_seq, n_csv = len(seq_paths), len(csv_paths)
if n_seq != n_csv:
raise IOError(
'Mismatch between image {} and annotation {} lengths'.format(
n_seq, n_csv))
class_names = open(class_names_path, 'r').readlines()
class_dict = {x.strip(): i for (i, x) in enumerate(class_names)}
print('class_dict: ', class_dict)
print('labels: ', labels)
if n_vis > 0:
if save:
save_fname = '{:s}_{:s}.{:s}'.format(save_dir, getDateTime(),
save_fmt)
save_path = os.path.join('log', save_fname)
writer = ImageWriter(save_path, _logger)
_logger.info('Saving {:s} image sequence to {:s}'.format(
save_fmt, save_path))
if n_seq % n_vis != 0:
raise AssertionError('n_seq: {} not multiple of n_vis: {}'.format(
n_seq, n_vis))
n_groups = int(n_seq / n_vis)
seq_id = 0
label = ''
for i in range(n_groups):
vis_gen = []
for j in range(n_vis):
if labels:
label = labels[j]
vis_gen.append(
visualize(params.vis,
_logger,
seq_paths[seq_id],
csv_paths[seq_id],
class_dict,
n_frames=n_frames,
generator_mode=1,
enable_masks=enable_masks,
label=label,
only_boxes=only_boxes,
crop_size=crop_size))
seq_id += 1
for imgs in zip(*vis_gen):
# img_stacked = np.hstack(imgs)
stack_params = {'grid_size': grid_size, 'preserve_order': 1}
if crop_size:
stack_params['annotations'] = labels
img_stacked = stackImages_ptf(imgs, **stack_params)
img_stacked = cv2.cvtColor(img_stacked, cv2.COLOR_RGB2BGR)
if vis_size:
img_stacked = resizeAR(img_stacked, vis_size[0],
vis_size[1])
if save:
writer.write(img_stacked)
if show_img:
cv2.imshow('img_stacked', img_stacked)
key = cv2.waitKey(1 - _pause) % 256
if key == 27:
break
elif key == ord('q'):
_quit = 1
break
elif key == 32:
_pause = 1 - _pause
if _quit:
break
if save:
writer.release()
else:
for i in range(n_seq):
visualize(params.vis,
_logger,
seq_paths[i],
csv_paths[i],
class_dict,
n_frames=n_frames,
enable_masks=enable_masks)
if _quit:
break
if status == "MATCH!":
cv2.rectangle(img, (int(bb[0]), int(bb[1])), (int(bb[2]), int(bb[3])), green, 2)
else:
cv2.rectangle(img, (int(bb[0]), int(bb[1])), (int(bb[2]), int(bb[3])), light_red, 2)
color = light_red
if status == "MATCH!":
color = green
text = "Result: " + status + " "
img, line_width = utils.draw_text_in_image(img, text, (margin + line_width, v_pos), color,
line_width)
if ovmax > 0: # if there is intersections between the bounding-boxes
bbgt = [float(x) for x in gt_match["bbox"].split()]
cv2.rectangle(img, (int(bbgt[0]), int(bbgt[1])), (int(bbgt[2]), int(bbgt[3])), light_blue, 2)
img_res_ar = resizeAR(img, save_w, save_h)
cv2.imshow("Animation", img_res_ar)
video_out.write(img_res_ar)
k = cv2.waitKey(1 - _pause)
if k == ord('q') or k == 27:
video_out.release()
cv2.destroyAllWindows()
sys.exit(0)
elif k == ord('c'):
end_class = 1
break
elif k == ord('s'):
break
elif k == 32:
_pause = 1 - _pause
bkg_pkl_path = os.path.join(bkg_path, 'bkg_imgs.pkl')
if load_bkg and os.path.isfile(bkg_pkl_path):
print('Loading background images from {}'.format(bkg_pkl_path))
with open(bkg_pkl_path, 'rb') as f:
bkg_imgs = pickle.load(f)
else:
print('Reading background image sequence from {}'.format(bkg_path))
bkg_imgs = []
for i in range(n_bkgs):
_bkg_fname = bkg_file_list[i]
bkg_img_path = os.path.join(bkg_path, _bkg_fname)
bkg_img = cv2.imread(bkg_img_path)
orig_shape = bkg_img.shape
if bkg_size:
bkg_img, resize_factor, _, _ = resizeAR(bkg_img, bkg_width, bkg_height,
return_factors=True, add_border=False)
print('bkg_img.shape: ', bkg_img.shape)
bkg_imgs.append({
'name': _bkg_fname,
'path': bkg_img_path,
'image': bkg_img,
'resize_factor': resize_factor,
'orig_shape': orig_shape}
)
sys.stdout.write('\rDone {:d} frames'.format(i + 1))
sys.stdout.flush()
print()
print('Saving background images to {}'.format(bkg_pkl_path))
with open(bkg_pkl_path, 'wb') as f:
pickle.dump(bkg_imgs, f, pickle.HIGHEST_PROTOCOL)
# gt_file = os.path.join(pkl_files_path, seq_name, '{}_ground_truth.json'.format(file_id))
# ground_truth_data = json.load(open(gt_file))
# if gt_file not in loaded_json:
# loaded_json.add(gt_file)
# for obj in ground_truth_data:
# obj["used"] = False
if show_animation and missing_detections:
# print('\nfile_id: ', file_id)
# print('missing_detections:\n ', missing_detections)
for bb in missing_detections:
cv2.rectangle(img, (bb[0], bb[1]), (bb[2], bb[3]), magenta,
2)
img = resizeAR(img, save_w, save_h)
img = cv2.copyMakeBorder(img,
0,
bottom_border,
0,
0,
cv2.BORDER_CONSTANT,
value=BLACK)
height, _ = img.shape[:2]
v_pos = int(height - margin - (bottom_border / 2))
text = "{}: {} ".format(seq_name, ground_truth_img[0])
img, line_width = utils.draw_text_in_image(
img, text, (margin, v_pos), white, 0)
text = "Class [" + str(class_index + 1) + "/" + str(
n_classes) + "]: " + class_name + " "
img, line_width = utils.draw_text_in_image(