I change code in detectron/utils/vis.py
In function vis_keypoints:
change:
kp_mask = np.copy(img)
to:
kp_mask=img
change:
return kp_mask, keypoint_pixel_1, keypoint_pixel_2
to:
return kp_mask
live keypoints file:
#!/usr/bin/env python2
"""Perform inference on a single image or all images with a certain extension
(e.g., .jpg) in a folder.
"""
from collections import defaultdict
import cv2
import argparse
import glob
import logging
import os
import sys
import time
from caffe2.python import workspace
from detectron.core.config import assert_and_infer_cfg
from detectron.core.config import cfg
from detectron.core.config import merge_cfg_from_file
from detectron.utils.io import cache_url
from detectron.utils.logging import setup_logging
from detectron.utils.timer import Timer
import detectron.core.test_engine as infer_engine
import detectron.datasets.dummy_datasets as dummy_datasets
import detectron.utils.c2 as c2_utils
import detectron.utils.vis as vis_utils
import pycocotools.mask as mask_util
c2_utils.import_detectron_ops()
# OpenCL may be enabled by default in OpenCV3; disable it because it's not
# thread safe and causes unwanted GPU memory allocations.
cv2.ocl.setUseOpenCL(False)
class pp(object):
def __init__(self, keypoint, box):
self.keypoint = keypoint
self.box = box
def __str__(self):
return str(self.keypoint) + "\n" + str(self.box) + "\n"
def proce(
im, boxes, segms=None, keypoints=None, thresh=0.9, kp_thresh=2):
"""Constructs a numpy array with the detections visualized."""
if isinstance(boxes, list):
boxes, segms, keypoints, classes = vis_utils.convert_from_cls_format(
boxes, segms, keypoints)
people = []
for i in range(0, len(boxes)):
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
people.append(pp(keypoints[i], boxes[i]))
for i in range(0, len(people)):
print(i)
print(people[i])
def dvp(path):
name = path[path.rfind('/')+1:]
folder = path[:path.rfind('/')]
return folder, name
def main(path, config_yaml='e2e_keypoint_rcnn_R-101-FPN_1x.yaml', model_pkl='model_final.pkl'):
merge_cfg_from_file(config_yaml)
cfg.NUM_GPUS = 1
weights = cache_url(model_pkl, 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(weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
# Capture video from file
cap = cv2.VideoCapture(0)
ret=cap.isOpened()
fo, fn=dvp(path)
fo_cf, fn_cf=dvp(config_yaml)
k=0;
while ret & (k!=27):
# Capture frame-by-frame
ret, frame = cap.read()
if ret:
timers = defaultdict(Timer)
t = time.time()
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(model, frame, None, timers=timers)
print('Inference time: {:.3f}s'.format(time.time() - t))
print(type(frame))
vis_utils.vis_one_image_opencv(
frame, # BGR -> RGB for visualization
cls_boxes,
cls_segms,
cls_keyps,
thresh=0.7,
kp_thresh=2,
show_box=True,
dataset=dummy_coco_dataset,
show_class=True
)
cv2.imshow('frame',frame)
proce(frame, cls_boxes, cls_segms, cls_keyps)
k=cv2.waitKey(1)
cap.release()
if __name__ == '__main__':
workspace.GlobalInit(['caffe2', '--caffe2_log_level=0'])
main('IMG_2335.MOV', 'e2e_keypoint_rcnn_R-101-FPN_s1x.yaml', 'model_final.pkl')
window will turn to black sometime