#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2017 Takuma Yagi <tyagi@iis.u-tokyo.ac.jp>
#
# Distributed under terms of the MIT license.
from __future__ import print_function
from __future__ import division
from six.moves import range
import os
import time
import json
import argparse
import joblib
from box import Box
import numpy as np
import cv2
from utils.dataset import SceneDatasetForAnalysis
from utils.plot import draw_line, draw_dotted_line, draw_x
from mllogger import MLLogger
logger = MLLogger(init=False)
if __name__ == "__main__":
"""
Write prediction as video from prediction.json
By default, blue is input, red is gt, green is prediction
This method requires vid, frame, pid but if the latter two is unspecified it'll give you candidates
Input:
* fold_path: fold id (e.g. 5fold_yymmss_HHMMSS/yymmss_HHMMSS)
* traj_type: trajectory type to output (0: front, 1: back, 2: across, 3:other)
* img_ratio: output image ratio (default: 0.5, 640x480)
* no_pred: do not plot prediction
* debug: if True, do not predict any images
* output_dir: directory to plot images and videos
* vid: video id (e.g. GOPR0235U20)
* frame: video frame number
* pid: person id number
Output:
plot images and videos to <output_dir>/yymmss_HHMMSS/
"""
parser = argparse.ArgumentParser()
# Data
parser.add_argument('fold_path', type=str)
parser.add_argument('--traj_type', type=int, default=-1)
parser.add_argument('--vid', type=str)
parser.add_argument('--frame', type=int, default=-1)
parser.add_argument('--pid', type=int, default=-1)
parser.add_argument('--nb_pause', type=int, default=5)
parser.add_argument('--img_ratio', type=float, default=0.5)
parser.add_argument('--ratio', type=float, default=0.7)
parser.add_argument('--thick', type=int, default=10)
parser.add_argument('--size_x', type=int, default=12)
parser.add_argument('--size_circle', type=int, default=20)
parser.add_argument('--fps', type=float, default=3)
parser.add_argument('--no_video', action='store_true')
parser.add_argument('--no_past', action='store_true')
parser.add_argument('--no_pred', action='store_true')
parser.add_argument('--no_gt', action='store_true')
parser.add_argument('--simultaneous', action='store_true')
# Others
parser.add_argument('--debug', action='store_true')
parser.add_argument('--output_dir', type=str, default="video_plots")
args = parser.parse_args()
start = time.time()
data_dir = "data"
logger.initialize(args.output_dir, debug=args.debug)
logger.info(vars(args))
save_dir = logger.get_savedir()
logger.info("Written to {}".format(save_dir))
nb_pause = args.nb_pause
output_image_list = []
prediction_path = os.path.join(args.fold_path, "prediction.json")
with open(prediction_path) as f:
data = json.load(f)
predictions = data["predictions"]
sargs = Box(data["arguments"])
print(predictions.keys())
if args.vid not in predictions:
print("No prediction found: aborted {} {} {}".format(args.vid, args.frame, args.pid))
exit(1)
if str(args.frame) not in predictions[args.vid] or str(args.pid) not in predictions[args.vid][str(args.frame)]:
print("Specify all vid, frame and pid: aborted {} {} {}".format(args.vid, args.frame, args.pid))
if args.vid in predictions and str(args.frame) in predictions[args.vid]:
print("Valid pids: {}".format(sorted(predictions[args.vid][str(args.frame)].keys())))
elif args.vid in predictions:
print("Valid frames {}".format(sorted(predictions[args.vid].keys())))
exit(1)
print("Loading data...")
dataset = joblib.load(sargs.in_data)
valid_split = sargs.eval_split + sargs.nb_splits * 2
valid_dataset = SceneDatasetForAnalysis(
dataset, sargs.input_len, sargs.offset_len if "offset" in sargs else 10, sargs.pred_len,
sargs.width, sargs.height)
color = (0, 255, 0)
color_ps = (255, 64, 64)
color_gt = (0, 0, 255)
color_pr = (0, 255, 0)
color_bb = (204, 102, 255)
# pid search
if args.frame != -1:
pid_list = []
for idx in range(len(valid_dataset)):
_, _, _, vid, frame, pid = valid_dataset.get_example(idx)[:6]
if vid == args.vid and frame == args.frame:
pid_list.append(pid)
print("Found pid: {}".format(np.unique(pid_list)))
if len(np.unique(pid_list)) == 1:
args.pid = pid_list[0]
print("Continue")
else:
exit(1)
idx_list = []
for idx in range(len(valid_dataset)):
_, _, _, vid, frame, pid = valid_dataset.get_example(idx)[:6]
if vid == args.vid and pid == args.pid:
idx_list.append((idx, frame))
if len(idx_list) == 0:
print("Not found: {} {}".format(args.vid, args.pid))
exit(1)
if args.frame == -1:
print("Candidate: {}".format(idx_list))
exit(1)
if args.frame not in map(lambda x: x[1], idx_list):
print("Wrong frame")
exit(1)
selected_idx = idx_list[list(map(lambda x: x[1], idx_list)).index(args.frame)][0]
past, ground_truth, pose, vid, frame, pid, flipped = valid_dataset.get_example(selected_idx)[:7]
img_dir = os.path.join(data_dir, "pseudo_viz", vid, "images")
def draw_past(img_dir, b, offset):
past, ground_truth, pose, vid, frame, pid, flipped = b[:7]
im = cv2.imread(os.path.join(img_dir, "rgb_{:05d}.jpg".format(int(frame) + offset)))
im = cv2.addWeighted(im, args.ratio, im, 0.0, 0)
# Past
if offset > 0:
cv2.circle(im, (int(past[0][0]), int(past[0][1])), args.size_circle, color_ps, -1)
im = draw_line(im, past[:offset+1], color_ps, 1.0, args.thick)
im = draw_x(im, past[offset], color_ps, 1.0, args.size_x)
return im
def draw_prediction(img_dir, b, offset_g, offset_p, predictions, img_offset, no_pred=False, no_gt=False):
past, ground_truth, pose, vid, frame, pid, flipped = b[:7]
vid, frame, pid, flipped, py, pe, pp, err, traj_type = predictions[vid][str(frame)][str(pid)]
predicted = np.array(py)[...,:2]
im = cv2.imread(os.path.join(img_dir, "rgb_{:05d}.jpg".format(int(frame) + img_offset)))
im = cv2.addWeighted(im, args.ratio, im, 0.0, 0)
cv2.circle(im, (int(past[0][0]), int(past[0][1])), args.size_circle, color_ps, -1)
im = draw_line(im, past, color_ps, 1.0, args.thick)
im = draw_x(im, past[-1], color_ps, 1.0, args.size_x)
if not args.no_gt:
if offset_g > 0:
im = draw_dotted_line(im, ground_truth[:offset_g+1], color_gt, 1.0, args.thick)
if offset_g > -1:
cv2.circle(im, (int(ground_truth[0][0]), int(ground_truth[0][1])), args.size_circle, color_gt, -1)
im = draw_x(im, ground_truth[offset_g], color_gt, 1.0, args.size_x)
if not no_pred:
if offset_p > 0:
im = draw_dotted_line(im, predicted[:offset_p+1], color_pr, 1.0, args.thick, 0.7)
if offset_p > -1:
im = draw_x(im, predicted[offset_p], color_pr, 1.0, args.size_x)
return im
def write_img(oframe, out_img):
cv2.imwrite("{}/{}_{}_{:02d}.jpg".format(save_dir, args.vid, args.pid, oframe), out_img)
start_b = valid_dataset.get_example(selected_idx)
input_len = sargs["input_len"]
pred_len = sargs["pred_len"]
# Input
if not args.no_past:
for x in range(input_len):
out = draw_past(img_dir, start_b, x)
write_img(x, out)
output_image_list.append(out)
if not args.no_pred and not args.simultaneous:
if not args.no_past:
# Pause
for x in range(nb_pause):
out = draw_past(img_dir, start_b, input_len-1)
output_image_list.append(out)
# Extend prediction
for x in range(pred_len):
out = draw_prediction(img_dir, start_b, -1, x, predictions, input_len - 1)
write_img(input_len + x, out)
output_image_list.append(out)
# Pause
for x in range(nb_pause):
out = draw_prediction(img_dir, start_b, -1, pred_len - 1, predictions, input_len - 1, args.no_pred, args.no_gt)
output_image_list.append(out)
# Actual future
for x in range(pred_len):
if args.simultaneous:
out = draw_prediction(img_dir, start_b, x, x, predictions, input_len + x, args.no_pred, args.no_gt)
else:
out = draw_prediction(img_dir, start_b, x, pred_len - 1, predictions, input_len + x, args.no_pred, args.no_gt)
write_img(input_len + pred_len + x, out)
output_image_list.append(out)
# Pause
for x in range(nb_pause):
out = draw_prediction(img_dir, start_b, pred_len - 1, pred_len - 1,
predictions, input_len + pred_len - 1, args.no_pred)
output_image_list.append(out)
if not args.no_video:
if args.no_pred:
video_name = os.path.join(save_dir, "{}_{}_{}_np.avi".format(args.vid, args.frame, args.pid))
#video_name = os.path.join(save_dir, "{}_{}_{}_np.mp4".format(args.vid, args.frame, args.pid))
else:
video_name = os.path.join(save_dir, "{}_{}_{}.avi".format(args.vid, args.frame, args.pid))
#video_name = os.path.join(save_dir, "{}_{}_{}.mp4".format(args.vid, args.frame, args.pid))
fourcc = cv2.VideoWriter_fourcc(*'XVID')
width, height = int(sargs.width * args.img_ratio), int(sargs.height * args.img_ratio)
writer = cv2.VideoWriter(video_name, fourcc, args.fps, (width, height))
#writer = cv2.VideoWriter(video_name, 0x21, args.fps, (width, height))
for img in output_image_list:
img = cv2.resize(img, (width, height))
writer.write(img)
writer.release()
logger.info("Wrote video into {}".format(video_name))
logger.info("Elapsed time: {} (s), Saved at {}".format(time.time()-start, save_dir))
