def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.NUM_GPUS = 1
# args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE)
assert_and_infer_cfg(cache_urls=False)
assert not cfg.MODEL.RPN_ONLY, \
'RPN models are not supported'
assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \
'Models that require precomputed proposals are not supported'
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
if os.path.isdir(args.im_or_folder):
im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext)
else:
im_list = [args.im_or_folder]
for i, im_name in enumerate(im_list):
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext)
)
logger.info('Processing {} -> {}'.format(im_name, out_name))
im = cv2.imread(im_name)
timers = defaultdict(Timer)
t = time.time()
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, im, None, timers=timers
)
logger.info('Inference time: {:.3f}s'.format(time.time() - t))
for k, v in timers.items():
logger.info(' | {}: {:.3f}s'.format(k, v.average_time))
if i == 0:
logger.info(
'\nNote: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)'
)
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=args.thresh,
kp_thresh=args.kp_thresh,
ext=args.output_ext,
out_when_no_box=args.out_when_no_box
)
# def vis_one_image_opencv(
# im, boxes, segms=None, keypoints=None, thresh=0.9, kp_thresh=2,
# show_box=False, dataset=None, show_class=False):
cvimg = vis_utils.vis_one_image_opencv(
im,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
show_class=True,
show_box=True,
thresh=args.thresh,
kp_thresh=args.kp_thresh
)
def vis(detection_dirs, output_dir, thresh, kp_thresh, average_frame, alpha,
save_as_images):
assert all([os.path.isdir(detect_dir) for detect_dir in detection_dirs
]), 'All input paths should be directories!'
detection_files = sorted(os.listdir(detection_dirs[0]))
assert all([
sorted(os.listdir(detect_dir)) == detection_files
for detect_dir in detection_dirs
])
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
for detect_file in detection_files:
all_detections = []
for detect_dir in detection_dirs:
with open(os.path.join(detect_dir, detect_file), 'rb') as f:
all_detections.append(pickle.load(f))
dummy_coco_dataset = dummy_datasets.get_coco_dataset_with_classes(
['__background__', 'car'])
boxes_seq, segms_seq, keyps_seq = [[
vid_detections[key] for vid_detections in all_detections
] for key in ['bboxes', 'segments', 'keypoints']]
vid_path = all_detections[0]['video_path']
orig_num_frms, orig_fps, out_fps, w, h = [
all_detections[0][kw] for kw in [
'original_num_frames', 'original_fps', 'out_fps', 'width',
'height'
]
]
time_stride = int(orig_fps / out_fps)
out_vid_path = os.path.join(output_dir, os.path.basename(vid_path))
vid_cap = cv2.VideoCapture(vid_path)
if save_as_images:
if not os.path.isdir(out_vid_path):
os.mkdir(out_vid_path)
else:
vid_writer = cv2.VideoWriter(out_vid_path, 1983148141, out_fps,
(w, h))
if average_frame:
vid_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
_, average_im = vid_cap.read()
first_step = time_stride
else:
first_step = 0
assert vid_cap.isOpened(), 'Cannot open input video.'
if not save_as_images:
assert vid_writer.isOpened(), 'Cannot create output video.'
for out_frm_id, orig_frm_id in tqdm(
enumerate(range(first_step, int(orig_num_frms), time_stride)),
total=int((orig_num_frms - first_step) // time_stride)):
vid_cap.set(cv2.CAP_PROP_POS_FRAMES, orig_frm_id)
_, im = vid_cap.read()
if average_frame:
average_im = np.around((1 - alpha) * average_im +
alpha * im).astype(np.uint8)
im = average_im
for i in range(len(all_detections)):
im = vis_utils.vis_one_image_opencv(
im,
boxes_seq[i][out_frm_id],
segms_seq[i][out_frm_id],
keyps_seq[i][out_frm_id],
dataset=dummy_coco_dataset,
show_class=True,
show_box=True,
thresh=thresh,
kp_thresh=kp_thresh,
)
if save_as_images:
cv2.imwrite(
os.path.join(out_vid_path,
str(out_frm_id) + '.jpg'), im)
else:
vid_writer.write(im)
vid_cap.release()
if not save_as_images:
vid_writer.release()
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.NUM_GPUS = 1
args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE)
assert_and_infer_cfg(cache_urls=False)
assert not cfg.MODEL.RPN_ONLY, \
'RPN models are not supported'
assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \
'Models that require precomputed proposals are not supported'
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
vid_cap = cv2.VideoCapture(args.video_fn)
width = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH) / args.downsize)
height = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT) / args.downsize)
#fourcc = cv2.VideoWriter_fourcc(*'XIVD')
#vid_out = cv2.VideoWriter('output.avi',-1, 20.0, (width,height))
FFMPEG_BIN = "ffmpeg"
command = [
FFMPEG_BIN,
'-y', # overwrite output
'-f',
'rawvideo',
'-vcodec',
'rawvideo',
'-s',
'{}x{}'.format(width, height),
'-pix_fmt',
'rgb24',
'-r',
'25', # fps
'-i',
'-', # input comes from a pipe
'-an', # no audio
'-vcodec',
'mpeg4',
'-b:v',
'30M',
'-maxrate',
'30M',
'-loglevel',
'panic',
'-nostdin',
args.output_fn
]
pipe = sp.Popen(command, stdin=sp.PIPE, stderr=sp.PIPE)
print(command)
ret = True
fid = 0
while ret:
print("time:", vid_cap.get(cv2.CAP_PROP_POS_MSEC) / 1000, "sec")
ret, im = vid_cap.read()
if not ret:
break
im = cv2.resize(im,
dsize=(width, height),
interpolation=cv2.INTER_LINEAR)
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, im, None, None)
rst_im = vis_utils.vis_one_image_opencv(im,
cls_boxes,
segms=cls_segms,
keypoints=cls_keyps,
thresh=args.thresh,
kp_thresh=args.kp_thresh,
show_box=True,
dataset=dummy_coco_dataset,
show_class=True)
rgb_im = cv2.cvtColor(rst_im, cv2.COLOR_BGR2RGB)
cv2.imshow('frame', rst_im)
pipe.stdin.write(rgb_im.tostring())
#pipe.stdin.flush()
if cv2.waitKey(1) & 0xFF == ord('q'):
break
vid_cap.release()
#Detect stuff and convert to usable form
cls_boxes, cls_segms, cls_keyps = kpdetection.detect(im)
boxes, segms, keyps, classes = \
vis_utils.convert_from_cls_format(cls_boxes,
cls_segms,
cls_keyps)
#Remove keypoints below thresholds
keyps, boxes = kpdetection.prune(keyps, boxes)
#Draw keypoints on frame
visualize = True
if visualize:
vis = vis_utils.vis_one_image_opencv(im,
cls_boxes,
keypoints=cls_keyps)
cv2.imshow("image", vis)
#SVM
if len(keyps) > 0: #if anything detected
instance = normalize_kp(keyps[0])
blackbox = np.zeros(shape=(255, 255), dtype=np.uint8)
for x in range(0, 17):
cv2.circle(blackbox, (instance[0][x], instance[1][x]),
1, (255, 255, 255),
thickness=2,
lineType=8,
shift=0)
cv2.imshow("keypoints", blackbox)
cv2.waitKey(1)