def load_model():
global file_num
file_num = 0
global class_model
class_model = feature_extractor()
model_path = './class_model.pki'
tmp = torch.load(model_path, map_location={'cuda:0': 'cpu'})
class_model.load_state_dict(tmp)
class_model.eval()
del tmp
global object_model
object_model = SSD(depth=50, width=1)
model_path = './ssd_patch.pki'
tmp = torch.load(model_path, map_location={'cuda:0': 'cpu'})
object_model.load_state_dict(tmp)
object_model.eval()
def initialize_net() -> None:
global ssd_net
# if already defined, return it
if ssd_net is not None:
print('use cached ssd_net')
return ssd_net
# get device ( cpu / gpu ) to be used
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print(f'device : {device}')
ssd_cfg = {
'num_classes':
num_classes, # number of classes including background class
'input_size': Parameters.IMG_SIZE,
'bbox_aspect_num': Parameters.BBOX_ASPECT_NUM,
'feature_maps': Parameters.FEATURE_MAPS,
'steps': Parameters.STEPS,
'min_sizes': Parameters.MIN_SIZES,
'max_sizes': Parameters.MAX_SIZES,
'aspect_ratios': Parameters.ASPECT_RATIOS,
'conf_thresh': Parameters.CONF_THRESHOLD,
'top_k': Parameters.TOP_K,
'nms_thresh': Parameters.NMS_THRESHOLD
}
print(f'initializing ssd with : {ssd_cfg}')
ssd_net = SSD(phase="inference", cfg=ssd_cfg)
# load weight created in training
weight_file_path = os.path.join(Parameters.ABEJA_TRAINING_RESULT_DIR,
'model.pth')
print(f'weight_file_path : {weight_file_path}')
# cf. https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-on-gpu-load-on-gpu
weight = torch.load(weight_file_path, map_location=device)
ssd_net.load_state_dict(weight)
ssd_net = ssd_net.to(device)
ssd_net.eval()
return ssd_net
def train(train_config):
logger = Logger(HOME+'/log', train_config.basenet)
if train_config.dataset_name == 'VOC':
cfg = voc_config
dataset = VOCDataset(DATA_DIR, transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif train_config.dataset_name == 'COCO':
cfg = coco_config
dataset = COCODataset(DATA_DIR, transform=SSDAugmentation(
cfg['min_dim'], MEANS))
if train_config.visdom:
import visdom
viz = visdom.Visdom()
ssd_net = SSD('train', train_config.basenet,
cfg['min_dim'], cfg['num_classes'], with_fpn=train_config.with_fpn)
net = ssd_net
if train_config.cuda:
net = nn.DataParallel(ssd_net)
cudnn.benchmark = True
if train_config.resume:
logger('Loading {} ...'.format(train_config.resume))
load_weights = torch.load(
train_config.resume, map_location=lambda storage, loc: storage)
ssd_net.load_state_dict(load_weights)
if train_config.cuda:
net = net.cuda()
if not train_config.resume:
logger('Initializing weights ...')
ssd_net.topnet.apply(weights_init)
ssd_net.loc_layers.apply(weights_init)
ssd_net.conf_layers.apply(weights_init)
optimizer = optim.Adam(net.parameters(), lr=train_config.lr,
weight_decay=train_config.weight_decay)
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, train_config.cuda)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
logger('Loading the dataset...')
epoch_size = len(dataset) // train_config.batch_size
logger('Training SSD on:{}'.format(dataset.name))
# logger('using the specified args:')
step_index = 0
if train_config.visdom:
vis_title = 'SSD.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
data_loader = data.DataLoader(dataset, train_config.batch_size,
num_workers=train_config.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
t0 = time.time()
for iteration in range(train_config.start_iter, cfg['max_iter']):
if train_config.visdom and iteration != 0 and (iteration % epoch_size == 0):
update_vis_plot(epoch, loc_loss.item(), conf_loss.item(), epoch_plot, None,
'append', epoch_size)
logger('epoch = {} : loss = {}, loc_loss = {}, conf_loss = {}'.format(
epoch, loc_loss + conf_loss, loc_loss, conf_loss))
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, train_config.lr,
train_config.gamma, step_index)
# load train data
images, targets = next(batch_iterator)
if iteration//epoch_size > 0 and iteration % epoch_size == 0:
batch_iterator = iter(data_loader)
print(iteration)
if train_config.cuda:
images = images.cuda()
targets = [ann.cuda()for ann in targets]
# else:
# images=torch.tensor(images)
# targets=torch.tensor(targets)
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
if train_config.visdom:
loc_loss += loss_l.item()
conf_loss += loss_c.item()
if iteration % 50 == 0:
t1 = time.time()
logger('timer: %.4f sec. || ' % (t1 - t0)+'iter ' + repr(iteration) +
' || Loss: %.4f ||' % (loss.item()) +
' || loc_loss: %.4f ||' % (loss_l.item()) +
' || conf_loss: %.4f ||' % (loss_c.item()))
t0 = time.time()
if train_config.visdom:
update_vis_plot(iteration, loss_l.item(), loss_c.item(),
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
logger('Saving state, iter:%d' % iteration)
torch.save(ssd_net.state_dict(), train_config.save_folder +
'ssd224_VOC_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
train_config.save_folder + 'ssd224_VOC.pth')
'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
}
#SSDネットワークモデル
net = SSD(phase="inference", cfg=ssd_cfg)
#device = torch.device('cpu')
#SSDの学習済み重みを設定
#net_weights = torch.load('./weights/ssd300_600.pth', map_location=device)
#net_weights = torch.load('./weights/ssd300_50.pth', map_location={'cuda0': 'cpu'})
net_weights = torch.load('./weights/ssd300_1.pth', map_location={'cuda0': 'cpu'})
#net_weights = torch.load('./weights/ssd300_mAP_77.43_v2.pth', map_location={'cuda0': 'cpu'})
###net_weights = torch.load('./weights/ssd300_mAP_77.43_v2.pth', map_location=device)
net.load_state_dict(net_weights)
print('ネットワーク設定完了：学習済みの重みをロードしました')
#画像読み込み
image_file_path = "./data/cowboy-757575_640.jpg"
img = cv2.imread(image_file_path) #[高さ][幅][色BGR]
height, width, chanels = img.shape #画像サイズを取得
#元画像の表示
plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.show()
#前処理クラスの作成
color_mean = (104, 117, 123) #(BGR)の色の平均値
input_size = 300 #画像のinputサイズを300×300にする
