def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
video_frame_length = video_frame(args) #Yunhan Yang pass in a list to save all the frames
print ("The video has length " + str(video_frame_length) + " frames")
f_image_path = args.im_or_folder+ "frame%d.jpg" % 0
fr = cv2.imread(f_image_path, 0)
origin_width, origin_height = fr.shape[:2]
for x in range(0,video_frame_length):
im_list = [args.im_or_folder+"frame" + str(x) + ".jpg"] #Yunhan Yang have to save frame in real folder and then read in
#maybe need need double for loop for list of frames
for i, im_name in enumerate(im_list):
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(im_name) + '.pdf')
)
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)'
)
#Yunhan Yang edit
if not os.path.exists(args.im_or_folder+ "/video"):
os.makedirs(args.im_or_folder+ "/video")
#Yunhan Yang edit Detectron/lib/utils/vis.py add make result as jpg than pdf
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=0.7,
kp_thresh=2
)
def main(args):
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)
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
cap = cv2.VideoCapture(0)
while True:
ret, frame = cap.read()
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)
image = vis_utils.vis_one_image_opencv(np.array(frame),
cls_boxes,
cls_segms,
cls_keyps,
thresh=0.7,
kp_thresh=2,
show_box=True,
show_class=True)
cv2.imshow('camera', image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
print("Time:", time.time() - t)
cap.release()
cv2.destroyAllWindows()
def main():
args = parser.parse_args()
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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) + '.pdf')
#)
#logger.info('Processing {} -> {}'.format(im_name, out_name))
#im = cv2.imread(im_name)
hintDataURL = bottle.request.forms.get("hint")
hintDataURL = re.sub('^data:image/.+;base64,', '', hintDataURL)
hintDataURL = base64.urlsafe_b64decode(hintDataURL.encode("ascii"))
hintDataURL = np.fromstring(hintDataURL, dtype=np.uint8)
hintDataURL = cv2.imdecode(hintDataURL, -1)
hstr = str(np.random.randint(100, 999))
cv2.imwrite('record/' + hstr + '.hint.png', hintDataURL)
im = cv2.imread('record/' + hstr + '.hint.png')
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)
return 'ok'
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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) + '.pdf')
)
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)'
)
def create_panoptic_segmentation(img,
cls_boxes,
cls_segms,
cls_keyps,
thres=0.7):
boxes, segms, keypoints, classes = vis_utils.convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
dataset = dummy_datasets.get_coco_dataset()
ade_out = np.zeros(img.shape[:2], dtype="uint8")
coco_out = np.zeros(img.shape[:2], dtype="uint8")
inst_out = np.zeros(img.shape[:2], dtype="uint8")
if segms is not None:
masks = mask_util.decode(segms)
cnt = 1
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
if boxes[i, -1] < thres: # Score too low
continue
mask = masks[..., i]
mask = np.nonzero(mask)
class_name = dataset.classes[classes[i]]
ade_idx = ade20k_utils.category_to_idx(class_name)
if ade_idx is not None:
ade_out[mask] = ade_idx
coco_out[mask] = i
inst_out[mask] = cnt
cnt += 1
out = np.stack([ade_out, coco_out, inst_out], axis=-1)
return out
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg(cache_urls=False)
model = infer_engine.initialize_model_from_cfg()
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) + '.pdf'))
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)')
t1 = time.time()
im = vis_utils.vis_one_image_opencv(im,
cls_boxes,
segms=cls_segms,
keypoints=cls_keyps,
thresh=0.7,
kp_thresh=2,
show_box=True,
dataset=dummy_coco_dataset,
show_class=True)
# 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=0.7,
# kp_thresh=2
# )
t2 = time.time() - t1
print("vis time %f ms" % (t2 * 1000))
cv2.imwrite(
args.output_dir + '/' + im_name.split('/')[-1].split('.')[0] +
'.jpg', im)
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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) + '.pdf')
# )
# logger.info('Processing {} -> {}'.format(im_name, out_name))
#set webcam
cam = cv2.VideoCapture(0)
while True:
#Fetch image from camera
ret_val, im = cam.read()
#uncomment to resize image
#im = cv2.resize(im, (1200,1024))
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 0:
logger.info(
' \ Note: 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
"dummy_name",
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
cv2.destroyAllWindows()
def __init__(self, config, **kwargs):
kwargs['handlers'] = [(r'/analyse', AnalyseHandler)]
super(Application, self).__init__(**kwargs)
merge_cfg_from_file(config.get('DETECTRON', 'CONFIG'))
cfg.TEST.WEIGHTS = config.get('DETECTRON', 'WEIGHTS')
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
self.model = infer_engine.initialize_model_from_cfg()
self.dummy_coco_dataset = dummy_datasets.get_coco_dataset()
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
logger.info('about to assert')
assert_and_infer_cfg()
logger.info('About to initialise model')
model = infer_engine.initialize_model_from_cfg()
logger.info("Done initialising")
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) + '.png'))
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)')
segmented_images, classes, scores, segmented_binary_masks = vis_utils.segmented_images_in_original_image_size(
im,
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2)
found = False
for index, value in enumerate(segmented_images):
if classes[index] == args.class_label and not found:
logger.info('Writing output file to: {}'.format(str(i)))
bin_mask = vis_utils.vis_binary_mask(
im, segmented_binary_masks[index])
cv2.imwrite(out_name, value)
cv2.imwrite(out_name.rstrip(".png") + "bin.png", bin_mask)
found = True
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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]
"""
Add support for webcam
"""
# Set and get camera from OpenCV
cam = cv2.VideoCapture(0)
im_name = 'tmp_im'
while True:
# Fetch image from camera
_, im = cam.read()
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 0:
logger.info(
' \ Note: 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=0.7,
kp_thresh=2,
ext='jpg' # default is PDF, but we want JPG.
)
cv2.destroyAllWindows()
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()
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)
im_list = search(args.im_or_folder, args.image_ext)
else:
im_list = [args.im_or_folder]
for i, im_name in enumerate(im_list):
#for i, im_name in im_list:
#out_name = os.path.join(
# args.output_dir, '{}'.format(os.path.basename(im_name) + '.pdf')
#)
out_name = im_name.replace(args.im_or_folder, args.output_dir)
par_path = os.path.dirname(out_name)
if not os.path.exists(par_path):
os.makedirs(par_path)
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(
' \ Note: 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,
par_path,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2,
ext='png'
)
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)
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = (
dummy_datasets.get_vg3k_dataset()
if args.use_vg3k else dummy_datasets.get_coco_dataset())
if os.path.isdir(args.im_or_folder):
im_list = sorted(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) + '.npz')
)
logger.info('Processing {} -> {}'.format(im_name, out_name))
im = cv2.imread(im_name)
if im is None:
logger.info('Unable to read image, skipping.')
continue
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)'
)
boxes, segms, classes = convert(cls_boxes, cls_segms)
classes = np.array(classes, dtype=np.uint16)
resolution = segms[0]['size']
segms = np.array([x['counts'] for x in segms]) # Run-length encoding
valid = boxes[:, 4] >= args.thresh
if args.filter_classes:
valid &= np.isin(classes, all_classes)
boxes = boxes[valid].copy()
classes = classes[valid].copy()
segms = segms[valid].copy()
output_name = os.path.basename(im_name)
np.savez(args.output_dir + '/' + output_name, boxes=boxes, segments=segms, classes=classes, resolution=resolution)
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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]
test_ids = []
rles = []
for i, im_name in enumerate(im_list):
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(im_name) + '.pdf'))
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)')
new_test_ids, new_rles = vis_utils.make_submission(im[:, :, ::-1],
im_name,
cls_boxes,
cls_segms,
cls_keyps,
thresh=0.7)
test_ids.extend(new_test_ids)
rles.extend(new_rles)
import pandas as pd
sub = pd.DataFrame()
sub['ImageId'] = test_ids
sub['EncodedPixels'] = pd.Series(rles).apply(
lambda x: ' '.join(str(y) for y in x))
sub.to_csv(args.output_dir +
'/e2e_mask_rcnn_R-50-FPN_1x-lr3e-3-nuclei-6-new.csv',
index=False)
def main(args):
logger = logging.getLogger(__name__)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
cfg_orig = yaml.load(yaml.dump(cfg))
im = cv2.imread(args.im_file)
if args.rpn_pkl is not None:
proposal_boxes, _proposal_scores = get_rpn_box_proposals(im, args)
workspace.ResetWorkspace()
else:
proposal_boxes = None
cls_boxes, cls_segms, cls_keyps = None, None, None
for i in range(0, len(args.models_to_run), 2):
pkl = args.models_to_run[i]
yml = args.models_to_run[i + 1]
cfg.immutable(False)
merge_cfg_from_cfg(cfg_orig)
merge_cfg_from_file(yml)
if len(pkl) > 0:
weights_file = pkl
else:
weights_file = cfg.TEST.WEIGHTS
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = model_engine.initialize_model_from_cfg(weights_file)
with c2_utils.NamedCudaScope(0):
cls_boxes_, cls_segms_, cls_keyps_ = \
model_engine.im_detect_all(model, im, proposal_boxes)
cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes
cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms
cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps
workspace.ResetWorkspace()
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(args.im_file) + '.pdf')
)
logger.info('Processing {} -> {}'.format(args.im_file, out_name))
vis_utils.vis_one_image(
im[:, :, ::-1],
args.im_file,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
def vis(im_name, im, cls_boxes, cls_segms, cls_keyps):
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
out_dir = None
loaded = vis_utils.vis_one_image_opencv(im,
cls_boxes,
segms=cls_segms,
keypoints=cls_keyps,
thresh=0.9,
kp_thresh=2,
show_box=False,
dataset=None,
show_class=False)
misc.imsave("loaded.png", loaded)
def main(args):
logger = logging.getLogger(__name__)
args.cfg = '/detectron/configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml'
args.weights = '/data/workspace/fbdet/models/mask_rcnn_R_101_FPN_2x/model_final.pkl'
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)
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]
im_list = []
im_list.append('/data/workspace/fbdet/test_pic/11.jpg')
video = cv2.VideoCapture('/data/pic/valid_video_00.avi')
frame = 0
while (True):
if frame > 0:
break
# ret, im = video.read()
# if im is None or ret is None:
# print("video.read() fail || video.read() is end!")
# break
im = cv2.imread(im_list[0])
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))
print('-----------------------------', frame)
boxs_list = vis_utils.get_boxes_image(cls_boxes,
cls_segms,
cls_keyps,
thresh=0.7,
dataset=dummy_coco_dataset)
print(boxs_list)
print('-----------------------------')
for i in range(len(boxs_list)):
box = boxs_list[i]
drawBoxOnImg(im, box[1], box[2], box[3], box[4], 0, 0, frame)
frame += 1
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)
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
for root_dir_path_1, sub_dir_path_list_1, sub_file_path_list_1 in os.walk(
args.im_or_folder):
sub_dir_path_list_1 = sorted(sub_dir_path_list_1)
for i, sub_dir_path_1 in enumerate(sub_dir_path_list_1):
for root_dir_path_2, sub_dir_path_list_2, sub_file_path_list_2 in os.walk(
os.path.join(root_dir_path_1, sub_dir_path_1)):
sub_file_path_list_2 = sorted(sub_file_path_list_2)
out_file = open(
os.path.join(args.output_dir,
sub_dir_path_1 + "_Det_ffasta.txt"), "wb")
for img_idx, sub_file_path_2 in enumerate(
sub_file_path_list_2):
im = cv2.imread(
os.path.join(root_dir_path_2, sub_file_path_2))
timers = defaultdict(Timer)
t = time.time()
if (img_idx + 1) % 1000 == 0:
sys.stdout.write(
"\rFinish {} images\n".format(img_idx + 1))
sys.stdout.flush()
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, im, None, timers=timers)
if isinstance(cls_boxes, list):
cls_boxes, cls_segms, cls_keyps, classes = vis_utils.convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
if cls_boxes is None or cls_boxes.shape[0] == 0:
continue
obj_idx = 0
for cls_box, cls in zip(cls_boxes, classes):
if int(cls) != 3 and int(cls) != 6:
continue
out_file.write("{},{},{},{},{},{},{}\n".format(
img_idx + 1, obj_idx + 1, cls_box[0],
cls_box[1], cls_box[2] - cls_box[0],
cls_box[3] - cls_box[1], cls_box[4]))
obj_idx += 1
out_file.close()
print("Finish {} / {} of video sequences".format(
i + 1, len(sub_dir_path_list_1)))
break
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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) + '.pdf')
)
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(
' \ Note: 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=0.7,
kp_thresh=2
)
def __init__(self):
self.cfg = rospy.get_param('~cfg', 'e2e_mask_rcnn_R-101-FPN_2x.yaml')
self.wts = rospy.get_param('~wts', 'model_final.pkl')
self.train_wts = rospy.get_param('~train_wts', 'R-101.pkl')
self.confidence = rospy.get_param('~confidence', 0.9)
self.sub_img_topic = rospy.get_param('~sub_img_topic',
"image_rect_color")
self.frame_rate = rospy.get_param('~frame_rate', 30)
self.gpu_id = rospy.get_param('~gpu_id', 0)
self.mask_on = rospy.get_param('~mask_on', True)
merge_cfg_from_file(self.cfg)
cfg.NUM_GPUS = 2
cfg.MODEL.MASK_ON = True
cfg.TRAIN.WEIGHTS = self.train_wts
assert_and_infer_cfg()
utils.logging.setup_logging(__name__)
self.logger = logging.getLogger(__name__)
self.model = infer_engine.initialize_model_from_cfg(self.wts,
gpu_id=self.gpu_id)
self.model1 = infer_engine.initialize_model_from_cfg(self.wts,
gpu_id=1)
self.dummy_coco_dataset = dummy_datasets.get_coco_dataset()
self.pupil_subscriber = rospy.Subscriber('/scene/left/fit_point',
ImagePoint,
self.pupil_callback,
queue_size=1)
self.image_subscriber = rospy.Subscriber(self.sub_img_topic,
Image,
self.callback,
queue_size=1)
# self.compressed_image_subscriber = rospy.Subscriber(self.sub_img_topic+"/compressed", CompressedImage, self.compimgcallback, queue_size=1)
self.pub_bboxes_topic = rospy.resolve_name(
self.sub_img_topic) + '/bboxes'
print("Mask RCNN Initialized")
self.bboxes_publisher = rospy.Publisher(self.pub_bboxes_topic,
BBoxDetArray,
queue_size=1)
self.pub_img_topic = self.sub_img_topic + "_detection"
self.image_publisher = rospy.Publisher(self.pub_img_topic,
Image,
queue_size=20)
self.bridge = CvBridge()
self.last_detect = rospy.Time.now()
self.fit_point = []
def main():
logger = logging.getLogger(__name__)
merge_cfg_from_file('/detectron/e2e_mask_rcnn_R-101-FPN_2x.yaml')
cfg.NUM_GPUS = 4
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg('/detectron/models/model_final.pkl')
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
# cam = cv2.VideoCapture("rtsp://192.168.128.12:554/mpeg4cif")
# cam = cv2.VideoCapture("rtsp://192.168.128.11:554/av0_1")
# cam = cv2.VideoCapture("http://192.168.128.14/video.cgi")
n = 0
tmp_file_name = '/tmp/tmp.jpg'
im0 = 0
im1 = 0
while True:
# ret_val, im = cam.read()
# cv2.imwrite(tmp_file_name, im)
im = cv2.imread('/detectron/k2m30/img/' + str(n) + '.jpg')
timers = defaultdict(Timer)
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, im, None, timers=timers
)
data = vis_utils.vis_one_image_opencv(
im, # BGR -> RGB for visualization
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
# box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2,
# ext='png'
)
# time.sleep(0.1)
# if data == None:
# logger.info(cls_boxes)
#
# else:
n += 1
n = n % 1000
file_name = '/tmp/' + str(n) + '.jpg'
cv2.imwrite(file_name, data)
logger.info(str(n) + ' saved')
logger.info(os.path.getsize(file_name))
def predict(cfg_path, weights_path):
logger = logging.getLogger(__name__)
merge_cfg_from_file(cfg_path)
cfg.NUM_GPUS = 1
gpu_id = 0 #指定GPU
weights_path = cache_url(weights_path, cfg.DOWNLOAD_CACHE)
assert_and_infer_cfg(cache_urls=False)
model = infer_engine.initialize_model_from_cfg(weights_path, gpu_id)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
#dummy_coco_dataset = dummy_datasets.get_illbuild_class11_dataset()
print("model is ok")
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) + '.jpg'))
logger.info('Processing {} -> {}'.format(im_name, out_name))
im = cv2.imread(im_name)
timers = defaultdict(Timer)
t = time.time()
with c2_utils.NamedCudaScope(gpu_id):
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(
' \ Note: 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=0.7,
kp_thresh=2)
def __init__(self):
# Logging initialization
workspace.GlobalInit(['caffe2', '--caffe2_log_level=0'])
utils.logging.setup_logging(__name__)
self.logger = logging.getLogger(__name__)
self.logger.info("In init of Segmenter, logging setup")
# Loading conf
self.conf = config.get_config()
# Model initialization
merge_cfg_from_file(self.conf['detectron']['cfg'])
cfg.NUM_GPUS = 1
weights = cache_url(self.conf['detectron']['weights'],
cfg.DOWNLOAD_CACHE)
assert_and_infer_cfg()
self.model = infer_engine.initialize_model_from_cfg(weights)
self.dummy_coco_dataset = dummy_datasets.get_coco_dataset()
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)
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = (dummy_datasets.get_vg3k_dataset() if args.use_vg3k
else dummy_datasets.get_coco_dataset())
if args.im_or_folder == 'train0':
#im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext)
im_list = cPickle.load(open('train_list.pkl'))
pre_folder = '/mnt/storage/jialinwu/detectron/detectron/datasets/data/coco/coco_train2014'
im_list = im_list[:40000]
elif args.im_or_folder == 'train1':
im_list = cPickle.load(open('train_list.pkl'))
pre_folder = '/mnt/storage/jialinwu/detectron/detectron/datasets/data/coco/coco_train2014'
im_list = im_list[40000:]
elif args.im_or_folder == 'val':
im_list = cPickle.load(open('val_list.pkl'))
pre_folder = '/mnt/storage/jialinwu/detectron/detectron/datasets/data/coco/coco_val2014'
elif args.im_or_folder == 'test0':
im_list = cPickle.load(open('test_list.pkl'))
im_list = im_list[:40000]
pre_folder = '../ARNet/image_captioning/data/images/test2015'
elif args.im_or_folder == 'test1':
im_list = cPickle.load(open('test_list.pkl'))
im_list = im_list[40000:]
pre_folder = '../ARNet/image_captioning/data/images/test2015'
for i in tqdm(range(len(im_list))):
im_name = pre_folder + '/' + im_list[i]
if im_name[-4:] != '.jpg':
continue
im = cv2.imread(im_name)
timers = defaultdict(Timer)
imgid = int(im_list[i][:-4].split('_')[-1])
if args.im_or_folder == 'val':
save_name = '/mnt/storage/jialinwu/seg_every_thing/npz_features/coco_val2014/%d.npz' % imgid
else:
save_name = '/mnt/storage/jialinwu/seg_every_thing/npz_features/coco_train2014/%d.npz' % imgid
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, im, None, save_name, timers=timers)
'''
def __init__(self, project, config, MAX=1000, dataset_name="ade"):
self.project = project
self.config = config
self.MAX = MAX
if dataset_name == "ade":
self.dataset = dummy_datasets.get_ade_dataset()
else:
self.dataset = dummy_datasets.get_coco_dataset()
self.out_dir = "tmp/html/"
self.images_dir = os.path.join(self.out_dir, "images/")
self.outfile = os.path.join(
self.out_dir, "{}_{}.html".format(project, int(time.time())))
if not os.path.exists(self.images_dir):
os.makedirs(self.images_dir)
self.init_outfile()
def main(args):
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
cfg_orig = yaml.load(yaml.dump(cfg))
im = cv2.imread(args.im_file)
if args.rpn_pkl is not None:
proposal_boxes, _proposal_scores = get_rpn_box_proposals(im, args)
workspace.ResetWorkspace()
else:
proposal_boxes = None
cls_boxes, cls_segms, cls_keyps = None, None, None
for i in range(0, len(args.models_to_run), 2):
pkl = args.models_to_run[i]
yml = args.models_to_run[i + 1]
merge_cfg_from_cfg(cfg_orig)
merge_cfg_from_file(yml)
if len(pkl) > 0:
cfg.TEST.WEIGHTS = pkl
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = model_engine.initialize_model_from_cfg()
with c2_utils.NamedCudaScope(0):
cls_boxes_, cls_segms_, cls_keyps_ = \
model_engine.im_detect_all(model, im, proposal_boxes)
cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes
cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms
cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps
workspace.ResetWorkspace()
vis_utils.vis_one_image(
im[:, :, ::-1],
args.im_file,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
def __init__(self, project, config, MAX=20000, dataset_name="ade"):
self.project = project
self.config = config
self.MAX = MAX
if dataset_name == "ade":
self.dataset = dummy_datasets.get_ade_dataset()
print "hi"
else:
self.dataset = dummy_datasets.get_coco_dataset()
self.out_dir = "tmp/{}/classes_{}/".format(self.project, dataset_name)
self.images_dir = os.path.join(self.out_dir, "images/")
if not os.path.exists(self.images_dir):
os.makedirs(self.images_dir)
self.print_location(self.out_dir)
self.out_paths = {}
self.refresh_rate = 100
def detectron_vizualize(config,
frame: FrameType,
cls_boxes: bytes,
cls_segms: bytes,
cls_keyps: bytes) -> FrameType:
cls_boxes = pickle.loads(cls_boxes)
cls_segms = pickle.loads(cls_segms)
cls_keyps = pickle.loads(cls_keyps)
vis_im = vis_utils.vis_one_image_opencv(
frame[:, :, ::1], # BGR -> RGB for visualization
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_datasets.get_coco_dataset(),
show_class=True,
thresh=0.7,
kp_thresh=2)
return vis_im
def main(args):
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
cfg_orig = yaml.safe_load(yaml.dump(cfg))
im = cv2.imread(args.im_file)
if args.rpn_pkl is not None:
proposal_boxes, _proposal_scores = get_rpn_box_proposals(im, args)
workspace.ResetWorkspace()
else:
proposal_boxes = None
cls_boxes, cls_segms, cls_keyps = None, None, None
for i in range(0, len(args.models_to_run), 2):
pkl = args.models_to_run[i]
yml = args.models_to_run[i + 1]
merge_cfg_from_cfg(cfg_orig)
merge_cfg_from_file(yml)
if len(pkl) > 0:
cfg.TEST.WEIGHTS = pkl
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = model_engine.initialize_model_from_cfg()
with c2_utils.NamedCudaScope(0):
cls_boxes_, cls_segms_, cls_keyps_ = \
model_engine.im_detect_all(model, im, proposal_boxes)
cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes
cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms
cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps
workspace.ResetWorkspace()
vis_utils.vis_one_image(im[:, :, ::-1],
args.im_file,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2)
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
with open(args.ann_json) as f:
im_list = gen_images(args.image_dir, json.loads(f.read())["images"])
with open(args.output_file, "w") as f:
for i, im_name in enumerate(im_list):
logger.info('Processing ' + im_name)
im = cv2.imread(im_name)
if im is None:
continue
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(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)')
result = vis_utils.vis_one_image(im_name,
cls_boxes,
cls_segms,
dataset=dummy_coco_dataset)
f.write(json.dumps(result) + "\n")
def camera(cam, queue, width, height, fps, args):
global running
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
start_time = 0
count = 0
# class_names =[
# '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
# 'bus', 'train', 'truck']
# color_list=[[0,0,0],[255,0,0],[0,255,0],[0,0,255],[255,255,0],[0,255,255],[255,255,0],[255,0,255],[255,255,255]]
class_names = [
'__background__', u'人', u'自行车', u'车', u'摩托车', 'airplane', u'车',
'train', u'车'
]
color_list = [[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255],
[255, 255, 0], [0, 0, 255], [255, 255, 0], [255, 0, 255],
[0, 0, 255]]
cls_sel = [1, 2, 3, 4, 6, 8]
cls_thresh = [1, 0.6, 0.5, 0.8, 0.5, 0.9, 0.7, 0.9, 0.5]
if count == 0:
logger.info(
' \ Note: inference on the first image will be slower than the '
'rest (caches and auto-tuning need to warm up)')
capture = cv2.VideoCapture(cam)
capture.set(cv2.CAP_PROP_FRAME_WIDTH, width)
capture.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
# capture.set(cv2.CAP_PROP_FPS, fps)
# size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
#print(cv2.__version__)
while (running):
frame = {}
ret, im = capture.read()
count = count + 1
if count == 5:
start_time = time.time()
# timers = defaultdict(Timer)
# # detect one image
with c2_utils.NamedCudaScope(0):
cls_boxes, _, _ = infer_engine.im_detect_all(model,
im,
None,
timers=None)
demo_vis_one_imageboxes_opencv(im,
cls_boxes,
thresh=cls_thresh,
show_box=True,
dataset=dummy_coco_dataset,
show_class=True,
class_names=class_names,
color_list=color_list,
cls_sel=cls_sel,
queue=q,
frame=frame,
count=count,
start_time=start_time)
queue.put(frame)
dtype=np.uint16).reshape(240, 320)
d4d = np.uint8(dmap.astype(float) * 255 / 2**12 - 1)
d4d = cv2.cvtColor(d4d, cv2.COLOR_GRAY2RGB)
d4d = 255 - d4d
return dmap, d4d
s = 0
done = False
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()
model = infer_engine.initialize_model_from_cfg(args.weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
while not done:
key = cv2.waitKey(1) & 255
if key == 27:
print("\tESC key detected!")
done = True
elif chr(key) == 's': #screen capture
print("\ts key detected. Saving image {}".format(s))
cv2.imwrite("ex2_" + str(s) + '.png', rgb)
im = get_rgb()
dmap, d4d = get_depth()
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)
def main(args):
logger = logging.getLogger(__name__)
merge_cfg_from_file(args.cfg)
cfg.TEST.WEIGHTS = args.weights
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg()
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) + '.pdf'))
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(
' \ Note: 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=0.7,
# kp_thresh=2
# )
boxes, segms, keypoints, classes = convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
masks = mask_util.decode(segms)
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < 0.7:
continue
name = get_class_string(classes[i], score, dummy_coco_dataset)
try:
os.mkdir(os.path.join(args.output_dir, name), 0755)
except OSError as exc:
if exc.errno != errno.EEXIST:
raise
pass
logger.info(name + ' {:0.2f}'.format(score).lstrip('0'))
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
r = cv2.boundingRect(c)
classout_name = os.path.join(
args.output_dir, name,
'{}_{}'.format(i, os.path.basename(im_name)))
cv2.imwrite(classout_name, im[r[1]:r[1] + r[3],
r[0]:r[0] + r[2]])
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True, device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
if args.merge_pdfs and num_images > 1:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
