def __init__(self, args):
torch.manual_seed(317)
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
heads = {'hm': args.num_classes, # cen, tl, tr, bl, br
'reg': 2*args.num_classes,
'wh': 2*4,
'mid_point': 2 * args.num_classes, }
self.model = spinal_net.SpineNet(heads=heads,
pretrained=True,
down_ratio=args.down_ratio,
final_kernel=1,
head_conv=256)
self.num_classes = args.num_classes
self.decoder = decoder.DecDecoder(K=args.K, conf_thresh=args.conf_thresh)
self.dataset = {'spinal': BaseDataset}
def __init__(self):
torch.manual_seed(317)
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
heads = {
'hm': 1,
'reg': 2,
'wh': 2 * 4,
}
self.model_path = "model_last.pth"
self.model = spinal_net.SpineNet(heads=heads,
basename='resnet34',
pretrained=True,
down_ratio=4,
final_kernel=1,
head_conv=256)
self.down_ratio = 4
self.model_state = False
self.decoder = decoder.DecDecoder(K=100, conf_thresh=0.2)