def load_model(self):
# load the model
self.sst = build_sst('test', 900)
self.sst.load_state_dict(
torch.load(config['resume'], map_location='cpu'))
self.sst.eval()
def init(img1_path, img2_path, model_path, cuda):
print('start init >>>>>>>>>>>>>>')
if not os.path.exists(img1_path) or not os.path.exists(
img2_path) or not os.path.exists(model_path):
raise ValueError("input parameter not right")
CompareTools.cuda = cuda
print('load image...')
# load image
CompareTools.img1 = cv2.imread(img1_path)
CompareTools.img2 = cv2.imread(img2_path)
CompareTools.img1_convert = CompareTools.convert_image(
CompareTools.img1, CompareTools.cuda)
CompareTools.img2_convert = CompareTools.convert_image(
CompareTools.img2, CompareTools.cuda)
CompareTools.img = np.concatenate(
[CompareTools.img1, CompareTools.img2], axis=0)
CompareTools.img_org = np.copy(CompareTools.img)
print('load model...')
# load net
CompareTools.sst = build_sst('test', 900, CompareTools.cuda)
if cuda:
cudnn.benchmark = True
CompareTools.sst.load_state_dict(torch.load(model_path))
CompareTools.sst = CompareTools.sst.cuda()
else:
CompareTools.sst.load_state_dict(torch.load(model_path))
print('finish init <<<<<<<<<<<<')
def load_model(self):
# load the model
self.sst = build_sst('test', 900)
if self.cuda:
cudnn.benchmark = True
self.sst.load_state_dict(torch.load(config['resume']))
self.sst = self.sst.cuda()
else:
self.sst.load_state_dict(
torch.load(config['resume'], map_location='cpu'))
self.sst.eval()
def load_model(self):
# load the model
self.sst = build_sst('test', 900)
if self.cuda:
cudnn.benchmark = True
self.sst.load_state_dict(torch.load(config['resume']))
self.sst = self.sst.cuda()
else:
self.sst.load_state_dict(torch.load(config['resume'], map_location='cpu'))
for param in self.sst.parameters():
param.requires_grad = False
def __init__(self):
self.tracks = list()
self.max_drawing_track = TrackerConfig.max_draw_track_node
self.cuda = TrackerConfig.cuda
self.recorder = FeatureRecorder()
self.frame_index = 0
# load the model
self.sst = build_sst('test', 900)
if self.cuda:
cudnn.benchmark = True
self.sst.load_state_dict(torch.load(TrackerConfig.sst_model_path))
self.sst = self.sst.cuda()
else:
self.sst.load_state_dict(
torch.load(config['resume'], map_location='cpu'))
self.sst.eval()
from torch.autograd import Variable
import torch.utils.data as data
import numpy as np
import argparse
from data.ua import UATrainDataset
from config.config import config
from layer.sst import build_sst
from layer.sst_loss import SSTLoss
from utils.augmentations import SSJEvalAugment, collate_fn
import time
from utils.operation import show_circle, show_batch_circle_image
import cv2
# build the model
sst = build_sst('test', 900)
if config['cuda']:
cudnn.benchmark = True
sst.load_state_dict(torch.load(config['resume']))
sst = sst.cuda()
else:
sst.load_state_dict(torch.load(config['resume'], map_location='cpu'))
sst.eval()
dataset = UATrainDataset(
config['ua_image_root'], config['ua_detection_root'],
config['ua_ignore_root'],
SSJEvalAugment(config['sst_dim'], config['mean_pixel']))
data_loader = data.DataLoader(dataset,
config['batch_size'],
save_weights_iteration = config['save_weight_every_epoch_num'] * config['epoch_size']
else:
stepvalues = (90000, 95000)
save_weights_iteration = 5000
gamma = args.gamma
momentum = args.momentum
if args.tensorboard:
from tensorboardX import SummaryWriter
if not os.path.exists(config['log_folder']):
os.mkdir(config['log_folder'])
writer = SummaryWriter(log_dir=config['log_folder'])
sst_net = build_sst('train')
net = sst_net
if args.cuda:
net = torch.nn.DataParallel(sst_net)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
sst_net.load_weights(args.resume)
else:
vgg_weights = torch.load(args.basenet)
print('Loading the base network...')
sst_net.vgg.load_state_dict(vgg_weights)