def main():
# # FOR DATA PARSING
# parse_data("./data/VOC2007/","./data/VOC2012/", "./output/")
# # FOR TRAINING
# train(resume=True, seed=94743, opt_level='O1', resize_dims=(500,500))
# # FOR EVALUATION AND TESTING
model = SSD(21)
state_dict = torch.load("./output/checkpoint.pt", map_location=device)
model.load_state_dict(state_dict)
model = model.to(device)
# # FOR EVALUATION ON THE TESTING SET
# test(model, "ALLPT", resize_dims=(500,500))
# test(model, "11PT", resize_dims=(500,500))
# # FOR DETECTION ON A SINGLE IMAGE
# image = "/home/numan947/MyHome/AIMLDL/Re:SSD - An Implementation of Single Shot MultiBox Detector/data/VOC2012/JPEGImages/2011_001254.jpg"
# original_image = Image.open(image, mode='r')
# original_image = original_image.convert('RGB')
# detect_single_image(model, original_image, min_score=0.25, max_overlap=0.4, top_k=200, resize_dims=(500,500)).show()
# # FOR DEMO USING CAMERA
# camera(model, min_score=0.3, max_overlap=0.25, top_k=20, resize_dims=(500,500))
# # FOR GENERATING SHOW CASE
ds = PascalGeneratorDataset("./output/", split="test")
generate(model, ds, n=50)
def train(seed, resume=False, opt_level='O1', resize_dims=(300, 300)):
clear_cuda()
torch.manual_seed(seed)
data_folder = "./output/"
keep_difficult = True
n_classes = len(label_map)
other_checkpoint = data_folder+"oth_checkpoint.pt"
model_checkpoint = data_folder+"checkpoint.pt"
early_stopping = EarlyStopping(save_model_name=model_checkpoint, patience=3)
pin_memory = True
batch_size = 8
accumulation_factor = 6
iterations = 100000
workers = 4*torch.cuda.device_count()
lr = 1e-3
decay_lr_at = [80000]
decay_lr_to = 0.2
early_stopper_lr_decrease = 0.5
early_stopper_lr_decrease_count = 7
momentum = 0.9
weight_decay = 5e-3
grad_clip = None
torch.backends.cudnn.benchmark = True
optimizer = None
start_epoch = None
history = pd.DataFrame()
history.index.name="Epoch"
history_path = "./output/HISTORY"+"_SEED_{}".format(seed)+".csv"
model = SSD(n_classes)
biases = list()
not_biases = list()
for param_name, param in model.named_parameters():
if param.requires_grad:
if param_name.endswith('.bias'):
biases.append(param)
else:
not_biases.append(param)
# optimizer = RAdam(params=[{'params':biases,'lr':2*lr},{'params':not_biases}], lr=lr, weight_decay=weight_decay)
# optimizer = torch.optim.RMSprop(params=[{'params':biases,'lr':2*lr},{'params':not_biases}], lr=lr, weight_decay=weight_decay, momentum=momentum)
# optimizer = torch.optim.SGD(params=[{'params':biases,'lr':2*lr},{'params':not_biases}], lr=lr, weight_decay=weight_decay, momentum=momentum)
optimizer = apex.optimizers.FusedLAMB(params=[{'params':biases,'lr':2*lr},{'params':not_biases}], lr=lr, weight_decay=weight_decay)#, momentum=momentum)
model = model.to(device)
criterion = MultiBoxLoss(model.priors_cxcy).to(device)
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level, loss_scale=1.0)
start_epoch = 0
if resume and os.path.exists(other_checkpoint):
ocheckpt = torch.load(other_checkpoint, map_location=device)
optimizer.load_state_dict(ocheckpt['optimizer_state'])
start_epoch = ocheckpt['epoch'] + 1
lr = get_lr(optimizer)
history = pd.read_csv(history_path, index_col="Epoch")
model.load_state_dict(torch.load(model_checkpoint, map_location=device))
# adjust_learning_rate(optimizer, 0.1)
lr = get_lr(optimizer)
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
amp.load_state_dict(ocheckpt['amp_state'])
train_dataset = PascalDataset(data_folder, split="train", keep_difficult=keep_difficult, resize_dims=resize_dims)
train_len = int(0.85*len(train_dataset))
valid_len = len(train_dataset) - train_len
train_data, valid_data = torch.utils.data.dataset.random_split(train_dataset, [train_len, valid_len])
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, collate_fn=train_dataset.collate_fn,
num_workers=workers, pin_memory=pin_memory)
valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=batch_size, shuffle=False, collate_fn=train_dataset.collate_fn,
num_workers=workers, pin_memory=pin_memory)
total_epochs = iterations//(len(train_dataset)//32)
decay_lr_at = [it//(len(train_dataset)//32) for it in decay_lr_at]
total_epochs = 125
decay_lr_at = [100]
print("Training For: {}".format(total_epochs))
print("Decay LR at:", decay_lr_at)
for epoch in range(start_epoch, total_epochs):
if epoch in decay_lr_at:
adjust_learning_rate(optimizer, decay_lr_to)
lr*=decay_lr_to
print("Learning Rate Adjusted")
st = time.time()
print("EPOCH: {}/{} -- Current LR: {}".format(epoch+1, total_epochs, lr))
clear_cuda()
tl,ta = train_single_epoch(epoch, model, train_loader, optimizer, criterion, plot=False, clip_grad=grad_clip, accumulation_factor=accumulation_factor)
clear_cuda()
vl, va = evaluate(model, valid_loader, criterion)
print_epoch_stat(epoch, time.time()-st, history=history, train_loss=tl, valid_loss=vl)
early_stopping(tl, model)
other_state = {
'epoch': epoch,
'optimizer_state': optimizer.state_dict(),
'amp_state': amp.state_dict()
}
torch.save(other_state, other_checkpoint)
history.loc[epoch, "LR"] = lr
history.to_csv(history_path)
if early_stopping.early_stop:
if early_stopper_lr_decrease_count>0:
early_stopper_lr_decrease_count = early_stopper_lr_decrease_count-1
adjust_learning_rate(optimizer, early_stopper_lr_decrease)
early_stopping.early_stop = False
early_stopping.counter = 0
lr*=early_stopper_lr_decrease
print("Learning Rate Adjusted")
accumulation_factor*=2
else:
break
num_workers=8)
testset = ImageSet(opt, transform, is_train=False)
testloader = torch.utils.data.DataLoader(testset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=8)
print('## Data preparation finish ##')
print('## Building net : SSD300 ##')
net = SSD(opt)
# 是否加载之前保存的模型
if args.resume:
print(' # Resuming from checkpoint # ')
checkpoint = torch.load(opt.ckpt_path)
net.load_state_dict(checkpoint['net'])
best_loss = checkpoint['loss']
start_epoch = checkpoint['epoch']
# 加载预训练的模型
else:
print(' # Loading pretrained model # ')
net.load_state_dict(torch.load(opt.pretrained_model))
criterion = MultiBoxLoss()
if use_cuda:
net.cuda()
criterion.cuda()
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
def train(self, vis=False):
print("begin training....")
if not os.path.exists('weights'):
os.mkdir('weights')
# Device settings
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device1 = torch.device(
"cuda:1" if torch.cuda.is_available() else "cpu")
if self.multigpu:
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.model = nn.DataParallel(self.model)
self.model = self.model.to(device)
eval_model = SSD(num_classes=cfg.num_classes,
num_blocks=cfg.mbox,
top_k=cfg.top_k,
conf_thresh=cfg.conf_thresh,
nms_thresh=cfg.nms_thresh,
variance=cfg.variance)
eval_model = eval_model.to(device1)
for item in self.model.parameters():
print(item.requires_grad)
total_epoch = cfg.epoch
criterion = DEC_loss(num_classes=cfg.num_classes,
variances=cfg.variance,
device=device)
optimizer = optim.SGD(params=filter(lambda p: p.requires_grad,
self.model.parameters()),
lr=cfg.init_lr,
momentum=0.9,
weight_decay=cfg.weight_decay)
# scheduler = lr_scheduler.StepLR(optimizer, step_size=cfg.lr_decay_epoch, gamma=0.1)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=cfg.milestones,
gamma=0.1)
print('Loading Datasets...')
dsets = PASCALVOC(root=cfg.root,
image_sets=cfg.train_sets,
transform=transforms.DEC_transforms(
phase='train',
size=cfg.img_size,
mean=cfg.means,
std=cfg.std))
dsets_val = PASCALVOC(root=cfg.root,
image_sets=cfg.test_sets,
transform=transforms.DEC_transforms(
phase='val',
size=cfg.img_size,
mean=cfg.means,
std=cfg.std))
dset_loaders = torch.utils.data.DataLoader(
dsets,
cfg.batch_size,
num_workers=4,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
if vis:
viewDatasets_DEC(dset_loaders)
train_loss_dict = []
mAP_dict = []
for epoch in range(total_epoch):
print('Epoch {}/{}'.format(epoch, total_epoch - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
scheduler.step()
self.model.train()
running_loss = 0.0
for data in dset_loaders:
inputs, target = data
inputs = inputs.to(device)
target = [item.to(device) for item in target]
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = self.model(inputs, phase)
# backprop
loss_l, loss_c = criterion(outputs, target)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
running_loss += loss.item()
epoch_loss = running_loss / len(dsets)
print('{} Loss: {:.6}'.format(epoch, epoch_loss))
train_loss_dict.append(epoch_loss)
np.savetxt('train_loss.txt', train_loss_dict, fmt='%.6f')
if epoch % 5 == 0:
torch.save(
self.model.state_dict(),
os.path.join(
'weights', '{:d}_{:.4f}_model.pth'.format(
epoch, epoch_loss)))
torch.save(self.model.state_dict(),
os.path.join('weights', 'end_model.pth'))
else:
if epoch % 5 == 0:
model_dict = self.model.state_dict()
val_dict = {k[7:]: v for k, v in model_dict.items()}
eval_model.load_state_dict(val_dict)
maps = self.eval(device1, eval_model, dsets_val)
mAP_dict.append(maps)
np.savetxt('mAP.txt', mAP_dict, fmt='%.6f')
class DEC_Module(object):
def __init__(self, multigpu, resume):
self.model = SSD(num_classes=cfg.num_classes,
num_blocks=cfg.mbox,
top_k=cfg.top_k,
conf_thresh=cfg.conf_thresh,
nms_thresh=cfg.nms_thresh,
variance=cfg.variance)
if resume is not None:
print('Resuming training weights from {} ...'.format(resume))
resume_dict = torch.load(resume)
resume_dict_update = {}
for k in resume_dict:
if k.startswith('module') and not k.startswith('module_list'):
resume_dict_update[k[7:]] = resume_dict[k]
else:
resume_dict_update[k] = resume_dict[k]
self.model.load_state_dict(resume_dict_update)
else:
resnet = "resnet101"
print('Resuming weights from {} ...'.format(resnet))
pre_trained_dict = model_zoo.load_url(model_urls[resnet])
model_dict = self.model.state_dict()
updated_dict = {
k: v
for k, v in pre_trained_dict.items() if k in model_dict
}
model_dict.update(updated_dict)
self.model.load_state_dict(model_dict)
self.multigpu = multigpu
def train(self, vis=False):
print("begin training....")
if not os.path.exists('weights'):
os.mkdir('weights')
# Device settings
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device1 = torch.device(
"cuda:1" if torch.cuda.is_available() else "cpu")
if self.multigpu:
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.model = nn.DataParallel(self.model)
self.model = self.model.to(device)
eval_model = SSD(num_classes=cfg.num_classes,
num_blocks=cfg.mbox,
top_k=cfg.top_k,
conf_thresh=cfg.conf_thresh,
nms_thresh=cfg.nms_thresh,
variance=cfg.variance)
eval_model = eval_model.to(device1)
for item in self.model.parameters():
print(item.requires_grad)
total_epoch = cfg.epoch
criterion = DEC_loss(num_classes=cfg.num_classes,
variances=cfg.variance,
device=device)
optimizer = optim.SGD(params=filter(lambda p: p.requires_grad,
self.model.parameters()),
lr=cfg.init_lr,
momentum=0.9,
weight_decay=cfg.weight_decay)
# scheduler = lr_scheduler.StepLR(optimizer, step_size=cfg.lr_decay_epoch, gamma=0.1)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=cfg.milestones,
gamma=0.1)
print('Loading Datasets...')
dsets = PASCALVOC(root=cfg.root,
image_sets=cfg.train_sets,
transform=transforms.DEC_transforms(
phase='train',
size=cfg.img_size,
mean=cfg.means,
std=cfg.std))
dsets_val = PASCALVOC(root=cfg.root,
image_sets=cfg.test_sets,
transform=transforms.DEC_transforms(
phase='val',
size=cfg.img_size,
mean=cfg.means,
std=cfg.std))
dset_loaders = torch.utils.data.DataLoader(
dsets,
cfg.batch_size,
num_workers=4,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
if vis:
viewDatasets_DEC(dset_loaders)
train_loss_dict = []
mAP_dict = []
for epoch in range(total_epoch):
print('Epoch {}/{}'.format(epoch, total_epoch - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
scheduler.step()
self.model.train()
running_loss = 0.0
for data in dset_loaders:
inputs, target = data
inputs = inputs.to(device)
target = [item.to(device) for item in target]
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = self.model(inputs, phase)
# backprop
loss_l, loss_c = criterion(outputs, target)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
running_loss += loss.item()
epoch_loss = running_loss / len(dsets)
print('{} Loss: {:.6}'.format(epoch, epoch_loss))
train_loss_dict.append(epoch_loss)
np.savetxt('train_loss.txt', train_loss_dict, fmt='%.6f')
if epoch % 5 == 0:
torch.save(
self.model.state_dict(),
os.path.join(
'weights', '{:d}_{:.4f}_model.pth'.format(
epoch, epoch_loss)))
torch.save(self.model.state_dict(),
os.path.join('weights', 'end_model.pth'))
else:
if epoch % 5 == 0:
model_dict = self.model.state_dict()
val_dict = {k[7:]: v for k, v in model_dict.items()}
eval_model.load_state_dict(val_dict)
maps = self.eval(device1, eval_model, dsets_val)
mAP_dict.append(maps)
np.savetxt('mAP.txt', mAP_dict, fmt='%.6f')
def test(self):
print('testing, evaluation mode...')
self.model.eval()
print('loading data...')
dsets = PASCALVOC(root=cfg.root,
image_sets=cfg.test_sets,
transform=transforms.DEC_transforms(
phase='val',
size=cfg.img_size,
mean=cfg.means,
std=cfg.std))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.model = self.model.to(device)
num_imgs = len(dsets)
test_timer = Timer()
cv2.namedWindow('img')
for i in range(num_imgs):
print('testing {}...'.format(dsets.img_ids[i]))
img, target = dsets.__getitem__(i)
ori_img = cv2.imread(dsets._imgpath % dsets.img_ids[i])
h, w, c = ori_img.shape
x = img.unsqueeze(0)
x = x.to(device)
test_timer.tic()
detections = self.model(x, 'test')
detect_time = test_timer.toc(average=False)
print('test time: {}'.format(detect_time))
for j in range(1, detections.size(1)):
dets = detections[0, j, :]
mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
dets = torch.masked_select(dets, mask).view(-1, 5)
if dets.shape[0] == 0:
continue
if j:
boxes = dets[:, 1:]
boxes[:, 0] *= h
boxes[:, 1] *= w
boxes[:, 2] *= h
boxes[:, 3] *= w
scores = dets[:, 0].cpu().numpy()
for box, score in zip(boxes, scores):
y1, x1, y2, x2 = box
y1 = int(y1)
x1 = int(x1)
y2 = int(y2)
x2 = int(x2)
cv2.rectangle(ori_img, (x1, y1), (x2, y2), (0, 255, 0),
2, 2)
cv2.putText(ori_img,
cfg.VOC_CLASSES[int(j)] + "%.2f" % score,
(x1, y1 + 20), cv2.FONT_HERSHEY_SIMPLEX,
0.6, (255, 0, 255))
cv2.imshow('img', ori_img)
k = cv2.waitKey(0)
if k & 0xFF == ord('q'):
cv2.destroyAllWindows()
exit()
cv2.destroyAllWindows()
exit()
def eval(self, device, eval_model, dsets):
# print('evaluation mode...')
# self.model.eval()
eval_model.eval()
output_dir = cfg.output_dir
if not os.path.exists(output_dir):
os.mkdir(output_dir)
else:
shutil.rmtree(output_dir)
os.mkdir(output_dir)
num_imgs = len(dsets)
total_time = 0
det_file = os.path.join(output_dir, 'detections.pkl')
# print('Detecting bounding boxes...')
all_boxes = [[[] for _ in range(num_imgs)]
for _ in range(cfg.num_classes)]
_t = {'im_detect': Timer(), 'misc': Timer()}
for i in range(num_imgs):
img, target = dsets.__getitem__(i)
ori_img = cv2.imread(dsets._imgpath % dsets.img_ids[i])
h, w, c = ori_img.shape
x = img.unsqueeze(0)
x = x.to(device)
_t['im_detect'].tic()
detections = eval_model(x, 'test')
detect_time = _t['im_detect'].toc(average=False)
# ignore the background boxes
for j in range(1, detections.size(1)):
dets = detections[0, j, :]
mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
dets = torch.masked_select(dets, mask).view(-1, 5)
if dets.shape[0] == 0:
continue
boxes = dets[:, 1:]
boxes[:, 0] *= h
boxes[:, 1] *= w
boxes[:, 2] *= h
boxes[:, 3] *= w
scores = dets[:, 0].cpu().numpy()
cls_dets = np.hstack((boxes.cpu().numpy(), scores[:, np.newaxis])) \
.astype(np.float32, copy=False)
all_boxes[j][i] = cls_dets
# print('img-detect: {:d}/{:d} {:.3f}s'.format(i + 1, num_imgs, detect_time))
total_time += detect_time
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
f.close()
print('Saving the results...')
for cls_ind, cls in enumerate(cfg.labelmap):
# print('Writing {:s} VOC results file'.format(cls))
filename = dec_eval.get_voc_results_file_template('test', cls)
with open(filename, 'wt') as f:
for im_ind, index in enumerate(dsets.img_ids):
dets = all_boxes[cls_ind + 1][im_ind]
if dets == []:
continue
# the VOCdevkit expects 1-based indices
for k in range(dets.shape[0]):
f.write(
'{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
index[1], dets[k, -1], dets[k, 0] + 1,
dets[k, 1] + 1, dets[k, 2] + 1,
dets[k, 3] + 1))
# print('Evaluating detections....')
print('average time is {}'.format(float(total_time) / num_imgs))
maps = dec_eval.do_python_eval(output_dir=output_dir, use_07=True)
return maps
def eval_single(self):
print('evaluation mode...')
self.model.eval()
print('loading data...')
dsets = PASCALVOC(root=cfg.root,
image_sets=cfg.test_sets,
transform=transforms.DEC_transforms(
phase='val',
size=cfg.img_size,
mean=cfg.means,
std=cfg.std))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.model = self.model.to(device)
output_dir = cfg.output_dir
if not os.path.exists(output_dir):
os.mkdir(output_dir)
else:
shutil.rmtree(output_dir)
os.mkdir(output_dir)
num_imgs = len(dsets)
det_file = os.path.join(output_dir, 'detections.pkl')
print('Detecting bounding boxes...')
all_boxes = [[[] for _ in range(num_imgs)]
for _ in range(cfg.num_classes)]
_t = {'im_detect': Timer(), 'misc': Timer()}
total_time = 0
for i in range(num_imgs):
img, target = dsets.__getitem__(i)
ori_img = cv2.imread(dsets._imgpath % dsets.img_ids[i])
h, w, c = ori_img.shape
x = img.unsqueeze(0)
x = x.to(device)
_t['im_detect'].tic()
detections = self.model(x, 'test')
detect_time = _t['im_detect'].toc(average=False)
total_time += detect_time
# ignore the background boxes
for j in range(1, detections.size(1)):
dets = detections[0, j, :]
mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
dets = torch.masked_select(dets, mask).view(-1, 5)
if dets.shape[0] == 0:
continue
boxes = dets[:, 1:]
boxes[:, 0] *= h
boxes[:, 1] *= w
boxes[:, 2] *= h
boxes[:, 3] *= w
boxes[:, 0] = np.maximum(0., boxes[:, 0])
boxes[:, 1] = np.maximum(0., boxes[:, 1])
boxes[:, 2] = np.minimum(h, boxes[:, 2])
boxes[:, 3] = np.minimum(w, boxes[:, 3])
scores = dets[:, 0].cpu().numpy()
cls_dets = np.hstack((boxes.cpu().numpy(), scores[:, np.newaxis])) \
.astype(np.float32, copy=False)
all_boxes[j][i] = cls_dets
print('img-detect: {:d}/{:d} {:.3f}s'.format(
i + 1, num_imgs, detect_time))
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
f.close()
print('average time is {}'.format(float(total_time) / num_imgs))
print('Saving the results...')
for cls_ind, cls in enumerate(cfg.labelmap):
print('Writing {:s} VOC results file'.format(cls))
filename = dec_eval.get_voc_results_file_template('test', cls)
with open(filename, 'wt') as f:
for im_ind, index in enumerate(dsets.img_ids):
dets = all_boxes[cls_ind + 1][im_ind]
if dets == []:
continue
# the VOCdevkit expects 1-based indices
for k in range(dets.shape[0]):
f.write(
'{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
index[1], dets[k, -1], dets[k, 0] + 1,
dets[k, 1] + 1, dets[k, 2] + 1,
dets[k, 3] + 1))
print('Evaluating detections....')
dec_eval.do_python_eval(output_dir=output_dir, use_07=True)
# coding:utf-8
from __future__ import print_function
import torch
import torchvision.transforms as tfs
import torch.nn.functional as F
from model import SSD
from data import PriorBox
from config import opt
from PIL import Image, ImageDraw
net = SSD(opt)
net.load_state_dict(torch.load(opt.ckpt_path)['net'])
net.eval()
# 加载测试图片
img = Image.open('/home/j/MYSSD/pytorch-ssd-master/image/img1.jpg')
img1 = img.resize((300, 300))
transform = tfs.Compose([
tfs.ToTensor(),
tfs.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])
img1 = transform(img1)
# 前向传播
loc, conf = net(img1[None, :, :, :])
# 将数据转换格式
prior_box = PriorBox(opt)
cfg = {
"num_classes": 21, #VOC data include 20 class + 1 background class
"input_size": 300, #SSD300
"bbox_aspect_num": [4, 6, 6, 6, 4,
4], # Tỷ lệ khung hình cho source1->source6`
"feature_maps": [38, 19, 10, 5, 3, 1],
"steps": [8, 16, 32, 64, 100, 300], # Size of default box
"min_size": [30, 60, 111, 162, 213, 264], # Size of default box
"max_size": [60, 111, 162, 213, 264, 315], # Size of default box
"aspect_ratios": [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
}
net = SSD(phase="inference", cfg=cfg)
net_weights = torch.load("./data/weights/ssd300_100.pth",
map_location={"cuda:0": "cpu"})
net.load_state_dict(net_weights)
def show_predict(img_file_path):
img = cv2.imread(img_file_path)
color_mean = (104, 117, 123)
input_size = 300
transform = DataTransform(input_size, color_mean)
phase = "val"
img_tranformed, boxes, labels = transform(img, phase, "", "")
img_tensor = torch.from_numpy(img_tranformed[:, :,
(2, 1, 0)]).permute(2, 0, 1)
net.eval()
])
test_set = VOC_loader(root="D:\Data\\voc\\2007",
year='2007_test',
image_set='train',
phase='TEST',
transform=transform)
# 3.test loader
test_loader = data.DataLoader(dataset=test_set,
batch_size=1,
collate_fn=test_set.collate_fn)
# 4. model load
net = SSD().to(device)
net.load_state_dict(torch.load('./saves/ssd.1.pth.tar'))
net.eval()
# 5. test
with torch.no_grad():
for image, target in test_loader:
# 사실 target 은 필요 없둠
# print(image.size()) # image : torch.Size([1, 3, 300, 300])
# print(len(target)) # target : list ( torch.Size([object_n, 5]) ) len(list) = 1
# image
# 각각의 tensor 들을 gpu 에 올리는 부분
image = image.to(device)
from model import SSD
import data
#setting
parser = argparse.ArgumentParser(description='Single Shot Multi Detector')
opts.setting(parser)
opt = parser.parse_args()
#random_seed
PRNG = RandomState(opt.seed)
# model
model = SSD(opt.n_classes)
cfg = model.config
decoder = MakeBox(cfg)
model.load_state_dict(torch.load(opt.weight_path))
model.cuda()
cudnn.benchmark = True
#dataload
dataset = data.loader_test(cfg, opt.data_path, PRNG)
print('size of dataset:', len(dataset))
def test():
print('testing....')
pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels = [], [], [], [], []
for i in range(len(dataset)):
img, loc, label = dataset[i]
gt_bboxes.append(loc)