class yolact_prediction(object):
def __init__(self):
parser = argparse.ArgumentParser(description='YOLACT Predict in ROS')
parser.add_argument(
'--visual_top_k',
default=100,
type=int,
help='Further restrict the number of predictions to parse')
parser.add_argument('--traditional_nms',
default=False,
action='store_true',
help='Whether to use traditional nms.')
parser.add_argument('--hide_mask',
default=False,
action='store_true',
help='Whether to display masks')
parser.add_argument('--hide_bbox',
default=True,
action='store_true',
help='Whether to display bboxes')
parser.add_argument('--hide_score',
default=True,
action='store_true',
help='Whether to display scores')
parser.add_argument(
'--show_lincomb',
default=False,
action='store_true',
help='Whether to show the generating process of masks.')
parser.add_argument(
'--no_crop',
default=False,
action='store_true',
help='Do not crop output masks with the predicted bounding box.')
parser.add_argument('--real_time',
default=True,
action='store_true',
help='Show the detection results real-timely.')
parser.add_argument(
'--visual_thre',
default=0.3,
type=float,
help='Detections with a score under this threshold will be removed.'
)
self.args = parser.parse_args()
r = rospkg.RosPack()
self.bridge = CvBridge()
self.path = r.get_path('yolact_prediction')
model_name = "/src/weights/best_89.48_res101_custom_610000.pth"
strs = model_name.split('_')
config = strs[-3] + "_" + strs[-2] + "_config"
update_config(config)
print("Using " + config + " according to the trained_model.")
with torch.no_grad():
self.cuda = torch.cuda.is_available()
if self.cuda:
cudnn.benchmark = True
cudnn.fastest = True
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
self.net = Yolact()
self.net.load_weights(self.path + model_name, self.cuda)
print('Model loaded.')
if self.cuda:
self.net = self.net.cuda()
self.time_here = 0
self.frame_times = MovingAverage()
#### Publisher
self.rgb_pub = rospy.Publisher("Yolact_predict_img/",
Image,
queue_size=1)
image_sub = rospy.Subscriber("/camera/color/image_raw",
Image,
self.img_cb,
queue_size=1)
print("============ Ready ============")
def img_cb(self, rgb_data):
self.rgb_data = rgb_data
if self.rgb_data is not None:
cv_image = self.bridge.imgmsg_to_cv2(self.rgb_data, "bgr8")
predict_img = self.predict(cv_image)
self.rgb_pub.publish(self.bridge.cv2_to_imgmsg(
predict_img, "bgr8"))
self.rgb_data = None
def predict(self, img):
rgb_origin = img
img_numpy = img
img = torch.from_numpy(img.copy()).float()
img = img.cuda()
img_h, img_w = img.shape[0], img.shape[1]
img_trans = FastBaseTransform()(img.unsqueeze(0))
net_outs = self.net(img_trans)
nms_outs = NMS(net_outs, 0)
results = after_nms(nms_outs,
img_h,
img_w,
crop_masks=not self.args.no_crop,
visual_thre=self.args.visual_thre)
torch.cuda.synchronize()
temp = self.time_here
self.time_here = time.time()
self.frame_times.add(self.time_here - temp)
fps = 1 / self.frame_times.get_avg()
frame_numpy = draw_img(results,
img,
self.args,
class_color=True,
fps=fps)
return frame_numpy
def onShutdown(self):
rospy.loginfo("Shutdown.")
torch.cuda.empty_cache()
update_config(config)
print(f'\nUsing \'{config}\' according to the trained_model.\n')
with torch.no_grad():
if cuda:
cudnn.benchmark = True
cudnn.fastest = True
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
dataset = COCODetection(cfg.dataset.valid_images,
cfg.dataset.valid_info,
augmentation=BaseTransform())
net = Yolact()
# net=onnx.load("yolact.onnx")
# net= torch.jit.script(net)
net.load_weights('weights/' + args.trained_model, cuda)
net.eval()
if (args.onnx):
ONNX_util.save_yolact(net, dataset, "yolact.onnx")
print('\nModel loaded.\n')
if cuda:
net = net.cuda()
evaluate(net, dataset, args.max_num, False, args.cocoapi,
args.traditional_nms)
# Don't use the timer during training, there's a race condition with multiple GPUs.
timer.disable_all()
cuda = torch.cuda.is_available()
if cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
net = Yolact()
net.train()
if args.resume == 'latest':
weight = glob.glob('weights/latest*')[0]
net.load_weights(weight)
resume_step = int(weight.split('.pth')[0].split('_')[-1])
print(f'\nResume training with \'{weight}\'.\n')
elif args.resume and 'yolact' in args.resume:
net.load_weights('weights/' + args.resume)
resume_step = int(args.resume.split('.pth')[0].split('_')[-1])
print(f'\nResume training with \'{args.resume}\'.\n')
else:
net.init_weights(backbone_path='weights/' + cfg.backbone.path)
print('\nTraining from begining, weights initialized.\n')
optimizer = optim.SGD(net.parameters(),
lr=cfg.lr,
momentum=cfg.momentum,
weight_decay=cfg.decay)
criterion = Multi_Loss(num_classes=cfg.num_classes,
default=False,
action='store_true',
help='Show the detection results real-timely.')
parser.add_argument(
'--visual_thre',
default=0.3,
type=float,
help='Detections with a score under this threshold will be removed.')
args = parser.parse_args()
args.cfg = re.findall(r'res.+_[a-z]+', args.weight)[0]
cfg = get_config(args, mode='detect')
cuda = torch.cuda.is_available()
net = Yolact(cfg)
net.load_weights(cfg.weight, cuda)
net.eval()
print(f'Model loaded with {cfg.weight}.\n')
if cuda:
cudnn.benchmark = True
cudnn.fastest = True
net = net.cuda()
with torch.no_grad():
# detect images
if cfg.image is not None:
dataset = COCODetection(cfg, mode='detect')
data_loader = data.DataLoader(dataset,
1,
num_workers=4,
config = 'yolact_resnet50_config'
update_config(config)
print(f'\nUsing \'{config}\' according to the trained_model.\n')
with torch.no_grad():
cuda = torch.cuda.is_available()
if cuda:
cudnn.benchmark = True
cudnn.fastest = True
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
net = Yolact()
net.load_weights(args.trained_model)
net.eval()
print('Model loaded.\n')
if cuda:
net = net.cuda()
# detect images
if args.image is not None:
images = glob.glob(args.image + '/*.jpg')
num = len(images)
for i, one_img in enumerate(images):
img_name = one_img.split('/')[-1]
img_origin = torch.from_numpy(cv2.imread(one_img)).cuda().float()
img_h, img_w = img_origin.shape[0], img_origin.shape[1]
# Don't use the timer during training, there's a race condition with multiple GPUs.
timer.disable_all()
cuda = torch.cuda.is_available()
if cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
net = Yolact()
net.train()
if args.resume == 'latest':
weight = glob.glob('weights/latest*')[0]
net.load_weights(weight, cuda)
resume_step = int(weight.split('.pth')[0].split('_')[-1])
print(f'\nResume training with \'{weight}\'.\n')
elif args.resume and 'yolact' in args.resume:
net.load_weights('weights/' + args.resume, cuda)
resume_step = int(args.resume.split('.pth')[0].split('_')[-1])
print(f'\nResume training with \'{args.resume}\'.\n')
else:
net.init_weights(backbone_path='weights/' + cfg.backbone.path)
print('\nTraining from begining, weights initialized.\n')
optimizer = optim.SGD(net.parameters(),
lr=cfg.lr,
momentum=cfg.momentum,
weight_decay=cfg.decay)
criterion = Multi_Loss(num_classes=cfg.num_classes,
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
dataset = COCODetection(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
augmentation=SSDAugmentation())
net = Yolact()
net.train()
# Don't use the timer during training, there's a race condition with multiple GPUs.
timer.disable_all()
if args.resume is not None:
net.load_weights(args.resume)
resume_epoch = int(args.resume.split('.')[0].split('_')[2])
resume_iter = int(args.resume.split('.')[0].split('_')[3])
print(
f'\nResume training at epoch: {resume_epoch}, iteration: {resume_iter}.'
)
else:
net.init_weights(backbone_path='weights/' + cfg.backbone.path)
print('\nTraining from begining, weights initialized.')
optimizer = optim.SGD(net.parameters(),
lr=cfg.lr,
momentum=cfg.momentum,
weight_decay=cfg.decay)
criterion = Multi_Loss(
num_classes=cfg.num_classes,