def main():
configs = parse_configs()
# Re-produce results
if configs.seed is not None:
random.seed(configs.seed)
np.random.seed(configs.seed)
torch.manual_seed(configs.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
if configs.gpu_idx is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
if configs.dist_url == "env://" and configs.world_size == -1:
configs.world_size = int(os.environ["WORLD_SIZE"])
configs.distributed = configs.world_size > 1 or configs.multiprocessing_distributed
if configs.multiprocessing_distributed:
configs.world_size = configs.ngpus_per_node * configs.world_size
mp.spawn(main_worker, nprocs=configs.ngpus_per_node, args=(configs, ))
else:
main_worker(configs.gpu_idx, configs)
def main():
configs = parse_configs()
# Re-produce results
if configs.seed is not None:
random.seed(configs.seed)
np.random.seed(configs.seed)
torch.manual_seed(configs.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
if configs.gpu_idx is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
if configs.dist_url == "env://" and configs.world_size == -1:
configs.world_size = int(os.environ["WORLD_SIZE"])
configs.distributed = configs.world_size > 1 or configs.multiprocessing_distributed
if configs.multiprocessing_distributed:
# Since we have ngpus_per_node processes per node, the total world_size
# needs to be adjusted accordingly
configs.world_size = configs.ngpus_per_node * configs.world_size
# Use torch.multiprocessing.spawn to launch distributed processes: the
# main_worker process function
mp.spawn(main_worker, nprocs=configs.ngpus_per_node, args=(configs, ))
else:
# Simply call main_worker function
main_worker(configs.gpu_idx, configs)
def main():
configs = parse_configs()
if configs.gpu_idx is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
if configs.dist_url == "env://" and configs.world_size == -1:
configs.world_size = int(os.environ["WORLD_SIZE"])
configs.distributed = configs.world_size > 1 or configs.multiprocessing_distributed
if configs.multiprocessing_distributed:
configs.world_size = configs.ngpus_per_node * configs.world_size
mp.spawn(main_worker, nprocs=configs.ngpus_per_node, args=(configs, ))
else:
main_worker(configs.gpu_idx, configs)
cv2.imwrite(os.path.join(configs.frame_dir, '{:06d}.jpg'.format(frame_idx)), ploted_img)
frame_pred_infor = {
'seg': prediction_seg,
'ball': prediction_ball_final
}
queue_frames.append(frame_pred_infor)
frame_idx += 1
if configs.output_format == 'video':
output_video_path = os.path.join(configs.save_demo_dir, 'result.mp4')
cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -b 5000k -c:v mpeg4 {}'.format(
os.path.join(configs.frame_dir), output_video_path)
os.system(cmd_str)
def plot_detection(img, ball_pos, seg_img, events):
"""Show the predicted information in the image"""
img = cv2.addWeighted(img, 1., seg_img * 255, 0.3, 0)
img = cv2.circle(img, tuple(ball_pos), 5, (255, 0, 255), -1)
event_name = 'is bounce: {}, is net: {}'.format(events[0], events[1])
img = cv2.putText(img, event_name, (100, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
return img
if __name__ == '__main__':
configs = parse_configs()
demo(configs=configs)