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)
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()