def eval():
# load net
num_classes = len(labelmap) + 1 # +1 for background
root_dir = os.path.dirname(args.trained_model)
model_name = os.path.basename(args.trained_model)
# load data
if args.dataset == 'VOC':
dataset = VOCDetection(args.root_dir, [('2007', set_type)],
BaseTransform(args.input_size, dataset_mean),
VOCAnnotationTransform(),
only_test=True)
else:
dataset = COCODetection(args.root_dir, 'test-dev2017',
BaseTransform(args.input_size, dataset_mean),
COCOAnnotationTransform())
net = build_net('test', args.input_size, num_classes)
state_dict = torch.load(args.trained_model)
# create new OrderedDict that does not contain `module.`
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
if 'base' in name:
name = name.replace('base', 'vgg')
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
net.eval()
print('Finished loading model:{}'.format(args.trained_model))
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(root_dir, net, dataset, model_name, retest=False)
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
# with open(det_file, 'rb') as f:
# all_boxes = pickle.load(f)
# print('LOADED')
dataset.evaluate_detections(all_boxes, save_folder)
if __name__ == '__main__':
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd('test', cfg, args.use_pred_module) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
print(net)
# load data
dataset = COCODetection(args.dataset_root,
image_set='minival2014',
transform=BaseTransform(cfg['min_dim'], MEANS),
target_transform=COCOAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder, net, args.cuda, dataset,
BaseTransform(net.size, MEANS), args.top_k, 512,
thresh=args.confidence_threshold)
def train(args):
create_time = time.strftime('%Y%m%d_%H%M', time.localtime(time.time()))
save_folder_path = os.path.join(args.save_folder, create_time)
# n_classes = [20, 80][args.dataset == 'COCO']
# n_classes = 91
if not ((args.train_image_folder and args.val_image_folder)
or args.annotation):
print("train/val image folder and annotation should not be None")
return
train_dataset = COCODetection(
root=args.root,
image_set=args.train_image_folder,
annotation_json=args.annotation,
transform=SSDAugmentation(img_size=args.image_size),
# transform = BaseTransform(img_size = args.image_size),
target_transform=COCOAnnotationTransform())
train_dataloader = DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate)
val_dataset = COCODetection(
root=args.root,
image_set=args.val_image_folder,
annotation_json=args.annotation,
transform=BaseTransform(img_size=args.image_size),
target_transform=COCOAnnotationTransform())
n_classes = train_dataset.get_class_size() + 1
if args.class_map_path:
train_dataset.get_class_map(args.class_map_path)
if args.model == "mobilenetv2":
model = MobileNetv2(
n_classes=n_classes,
width_mult=args.width_mult,
round_nearest=8,
dropout_ratio=args.dropout_ratio,
use_batch_norm=True,
)
ssd = create_mobilenetv2_ssd_lite(model,
n_classes,
width_mult=args.width_mult,
use_batch_norm=True)
elif args.model == "mobilenetv3":
model = MobileNetv3(model_mode=args.model_mode,
n_classes=n_classes,
width_mult=args.width_mult,
dropout_ratio=args.dropout_ratio)
ssd = create_mobilenetv3_ssd_lite(model,
n_classes,
model_mode=args.model_mode)
else:
print("model structure only accept mobilenetv2 or mobilenetv3")
return
print("builded ssd module")
if GPU:
import torch.backends.cudnn as cudnn
model.cuda()
ssd.cuda()
cudnn.benchmark = True
if args.pretrain_model:
ssd.load_state_dict(
torch.load(args.pretrain_model, map_location=torch.device('cpu')))
elif args.pretrain_tfmodel and args.pretrain_tfmodel_weight_list:
ssd_state_dict = ssd.state_dict()
tf_weights_dict = load_tf_weights(args, ssd_state_dict)
ssd.load_state_dict(tf_weights_dict)
optimizer = optim.Adam(ssd.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(n_classes,
overlap_thresh=args.overlap_threshold,
prior_for_matching=True,
bkg_label=0,
neg_mining=True,
neg_pos=args.neg_pos_ratio,
neg_overlap=0.5,
encode_target=False)
with torch.no_grad():
if args.model == "mobilenetv2":
prior_box = PriorBox(MOBILEV2_300)
elif args.model == "mobilenetv3":
prior_box = PriorBox(MOBILEV3_300)
priors = Variable(prior_box.forward())
print("created default bbox")
n_train = min(train_dataset.__len__(), 5000)
n_val = min(val_dataset.__len__(), 1000)
global_step = 0
val_global_step = 0
writer = SummaryWriter(log_dir=args.summary_path)
for epoch in range(args.epochs):
mean_loss_conf = 0
mean_loss_loc = 0
inference_count = 0
ssd.train()
with tqdm(total=n_train,
desc=f"{epoch + 1} / {args.epochs}",
unit='img') as pbar:
for img, target in train_dataloader:
if GPU:
img = Variable(img.cuda())
target = [Variable(anno.cuda()) for anno in target]
else:
img = Variable(img)
target = [Variable(anno) for anno in target]
optimizer.zero_grad()
inference = ssd(img)
loss_loc, loss_conf = criterion(inference, priors, target)
writer.add_scalar('Train/location_loss', float(loss_loc),
global_step)
writer.add_scalar('Train/confidence_loss', float(loss_conf),
global_step)
pbar.set_postfix(
**{
"location loss": float(loss_loc),
"confidence loss": float(loss_conf)
})
mean_loss_loc += float(loss_loc)
mean_loss_conf += float(loss_conf)
total_loss = loss_loc + loss_conf
total_loss.backward()
# # clip gradient
# # clip_grad_norm_(net.parameters(), 0.1)
optimizer.step()
pbar.update(img.shape[0])
global_step += 1
inference_count += img.shape[0]
if inference_count > n_train: break
pbar.set_postfix(
**{
"location loss": float(mean_loss_loc / n_train),
"confidence loss": float(mean_loss_conf / n_train)
})
ssd.eval()
val_mean_loss_loc = 0
val_mean_loss_conf = 0
with tqdm(total=n_val, desc="Validation", unit="img") as vpbar:
for i in range(n_val):
img = val_dataset.get_image(i)
img = cv2.resize(img, (args.image_size, args.image_size))
height, width, _ = img.shape
target = val_dataset.get_annotation(i, width, height)
if GPU:
img = torch.from_numpy(
np.expand_dims(img.transpose(2, 0, 1),
0)).to(dtype=torch.float32).cuda()
target = torch.FloatTensor(target).unsqueeze(0).cuda()
else:
img = torch.from_numpy(
np.expand_dims(img.transpose(2, 0, 1),
0)).to(dtype=torch.float32)
target = torch.FloatTensor(target).unsqueeze(0)
inference = ssd(img)
loss_loc, loss_conf = criterion(inference, priors, target)
val_mean_loss_loc += float(loss_loc)
val_mean_loss_conf += float(loss_conf)
vpbar.set_postfix(
**{
'location loss': float(loss_loc),
'confidnece loss': float(loss_conf)
})
vpbar.update(1)
vpbar.set_postfix(
**{
'location loss': float(val_mean_loss_loc / n_val),
'confidnece loss': float(val_mean_loss_conf / n_val)
})
writer.add_scalar('Test/location_loss',
float(val_mean_loss_loc / n_val),
val_global_step)
writer.add_scalar('Test/confidence_loss',
float(val_mean_loss_conf / n_val),
val_global_step)
val_global_step += 1
if epoch % 10 == 0 or epoch == args.epochs - 1:
save_model(save_folder_path, ssd, epoch)
writer.close()
if module_path not in sys.path:
sys.path.append(module_path)
from utils.utils import color_list, vis_detections
import torch
from torch.autograd import Variable
import numpy as np
import cv2
from data import COCODetection, COCO_ROOT, COCOAnnotationTransform
from models.refinedetlite import build_refinedet
from data import COCO_CLASSES as labels
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net = build_refinedet('test', 320, 81) # initialize SSD
net.load_weights('../weights/RefineDetLiteCOCO/RefineDet320_COCO_138000.pth')
testset = COCODetection(COCO_ROOT, "val2017", None, COCOAnnotationTransform())
img_id = 121
image = testset.pull_image(img_id)
x = cv2.resize(image, (320, 320)).astype(np.float32)
x -= (104.0, 117.0, 123.0)
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
x = torch.from_numpy(x).permute(2, 0, 1)
xx = Variable(x.unsqueeze(0)) # wrap tensor in Variable
if torch.cuda.is_available():
xx = xx.cuda()
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
net.eval()
print('Finished loading model!')
# load data
if args.dataset == 'VOC':
dataset = VOCDetection(args.voc_root, [('0712', "2007_test")], None,
AnnotationTransform())
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, [('2014', 'minival')], None,
COCOAnnotationTransform())
#COCOroot, [('2015', 'test-dev')], None)
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
top_k = 200
save_folder = os.path.join(args.save_folder, args.dataset)
if args.version == "drf_refine_vgg":
detector = Detect(num_classes, 0, cfg, use_arm=True)
else:
detector = Detect(num_classes, 0, cfg)
test_net(save_folder,
net,
detector,
gamma = 0.1
momentum = 0.9
dataset_name = args.dataset
if dataset_name[0] == "V":
cfg = (VOC_300, VOC_512)[args.size == '512']
train_dataset = VOCDetection(VOCroot, datasets_dict[dataset_name],
SSDAugmentation(img_dim, bgr_means),
AnnotationTransform(), dataset_name)
# train_dataset = VOCDetection(VOCroot, datasets_dict[dataset_name], preproc(img_dim, bgr_means, p), AnnotationTransform())
test_dataset = VOCDetection(VOCroot, datasets_dict["VOC2007"], None,
AnnotationTransform(), dataset_name)
elif dataset_name[0] == "C":
train_dataset = COCODetection(COCOroot, datasets_dict[dataset_name],
SSDAugmentation(img_dim, bgr_means),
COCOAnnotationTransform(), dataset_name)
test_dataset = COCODetection(COCOroot, datasets_dict["COCOval"], None,
COCOAnnotationTransform(), dataset_name)
cfg = (COCO_300, COCO_512)[args.size == '512']
else:
print('Unkown dataset!')
if args.version == "ssd_vgg":
from models.ssd.vgg_net import build_ssd
print("ssd vgg")
elif args.version == "ssd_res":
from models.ssd.res_net import build_ssd
print("ssd resnet")
elif args.version == "drf_ssd_vgg":
from models.drfssd.vgg_drfnet import build_ssd
print("drf ssd vgg")