def __init__(self, load_path):
self.img = None
print('Loading model..')
self.net = DSOD(num_classes=10)
self.net.cuda()
self.net = torch.nn.DataParallel(self.net,
device_ids=range(
torch.cuda.device_count()))
self.net.load_state_dict(torch.load(load_path)['net'], False)
self.net.eval()
self.box_coder = SSDBoxCoder(self.net.module)
def test_eval():
net = DSOD(num_classes=21)
net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(torch.cuda.device_count()))
net.load_state_dict(torch.load(opt.load_path)['net'])
aps = eval(net)
print(aps['ap'])
print(aps['map'])
def main(**kwargs):
opt._parse(kwargs)
vis = Visualizer(env=opt.env)
# Model
print('==> Building model..')
net = DSOD(num_classes=21)
# Dataset
print('==> Preparing dataset..')
box_coder = SSDBoxCoder(net)
trainset = ListDataset(root=opt.train_img_root,
list_file=opt.train_img_list,
transform=Transform(box_coder, True))
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_worker,
pin_memory=True)
net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
criterion = SSDLoss(num_classes=21)
optimizer = optim.SGD(net.parameters(),
lr=opt.lr,
momentum=0.9,
weight_decay=1e-4)
best_map_ = 0
best_loss = 1e100
start_epoch = 0
if opt.load_path is not None:
print('==> Resuming from checkpoint..')
checkpoint = torch.load(opt.load_path)
net.load_state_dict(checkpoint['net'])
best_loss = checkpoint['map']
start_epoch = checkpoint['epoch'] + 1
print('start_epoch = ', start_epoch, 'best_loss = ', best_loss)
for epoch in range(start_epoch, start_epoch + 100):
print('\nEpoch: ', epoch)
net.train()
train_loss = 0
optimizer.zero_grad()
ix = 0
for batch_idx, (inputs, loc_targets,
cls_targets) in tqdm(enumerate(trainloader)):
inputs = Variable(inputs.cuda())
loc_targets = Variable(loc_targets.cuda())
cls_targets = Variable(cls_targets.cuda())
loc_preds, cls_preds = net(inputs)
ix += 1
loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets)
loss.backward()
train_loss += loss.data.item()
current_loss = train_loss / (1 + batch_idx)
if (batch_idx + 1) % (opt.iter_size) == 0:
for name, p in net.named_parameters():
p.grad.data.div_(ix)
ix = 0
optimizer.step()
optimizer.zero_grad()
if (batch_idx + 1) % opt.plot_every == 0:
vis.plot('loss', current_loss)
# img = predict(net, box_coder, os.path.join(opt.train_img_root, trainset.fnames[batch_idx]))
# vis.img('predict', np.array(img).transpose(2, 0, 1))
# if os.path.exists(opt.debug_file):
# import ipdb
# ipdb.set_trace()
print('current_loss: ', current_loss, 'best_loss: ', best_loss)
if (epoch + 1) % 20 == 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.1
if (epoch + 1) % opt.save_state_every == 0:
state = {
'net': net.state_dict(),
'map': current_loss,
'epoch': epoch,
}
torch.save(state, opt.checkpoint + '%s.pth' % epoch)
if current_loss < best_loss:
best_loss = current_loss
print('saving model at epoch: ', epoch)
state = {
'net': net.state_dict(),
'map': best_loss,
'epoch': epoch,
}
torch.save(state, opt.checkpoint + 'dsod.pth')
from __future__ import print_function
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image, ImageDraw
from torch.autograd import Variable
from torchcv.models import DSOD, SSDBoxCoder
from torchcv.datasets import ListDataset
from torchcv.transform import resize
from torchcv.utils.config import opt
import random
print('Loading model..')
net = DSOD(num_classes=21)
net.cuda()
net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
net.load_state_dict(torch.load(opt.load_path)['net'])
net.eval()
print('Loading image..')
def caffe_normalize(x):
return transforms.Compose([
transforms.Lambda(lambda x:255*x[[2,1,0]]) ,
transforms.Normalize([104,117,123], (1,1,1)), # make it the same as caffe
# bgr and 0-255
])(x)
def transform(img, boxes, labels):
img, boxes = resize(img, boxes, size=(opt.img_size, opt.img_size))
img = transforms.Compose([
def main(**kwargs):
opt._parse(kwargs)
vis = Visualizer(env=opt.env)
# Model
print('==> Building model..')
net = DSOD(num_classes=21)
start_epoch = 0 # start from epoch 0 or last epoch
if opt.load_path is not None:
print('==> Resuming from checkpoint..')
checkpoint = torch.load(opt.load_path)
net.load_state_dict(checkpoint['net'])
# Dataset
print('==> Preparing dataset..')
box_coder = SSDBoxCoder(net)
trainset = ListDataset(root=opt.data_root,
list_file=[opt.voc07_trainval, opt.voc12_trainval],
transform=Transform(box_coder, True))
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=8)
net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
criterion = SSDLoss(num_classes=21)
optimizer = optim.SGD(net.parameters(),
lr=opt.lr,
momentum=0.9,
weight_decay=5e-4)
best_map_ = 0
for epoch in range(start_epoch, start_epoch + 200):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
for batch_idx, (inputs, loc_targets,
cls_targets) in tqdm(enumerate(trainloader)):
inputs = Variable(inputs.cuda())
loc_targets = Variable(loc_targets.cuda())
cls_targets = Variable(cls_targets.cuda())
optimizer.zero_grad()
loc_preds, cls_preds = net(inputs)
loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets)
loss.backward()
optimizer.step()
train_loss += loss.data[0]
if (batch_idx + 1) % opt.plot_every == 0:
vis.plot('loss', train_loss / (batch_idx + 1))
img = predict(
net, box_coder,
os.path.join(opt.data_root, trainset.fnames[batch_idx]))
vis.img('predict', np.array(img).transpose(2, 0, 1))
if os.path.exists(opt.debug_file):
import ipdb
ipdb.set_trace()
aps = eval(net.module, test_num=epoch * 100 + 100)
map_ = aps['map']
if map_ > best_map_:
print('Saving..')
state = {
'net': net.state_dict(),
'map': best_map_,
'epoch': epoch,
}
best_map_ = map_
if not os.path.isdir(os.path.dirname(opt.checkpoint)):
os.mkdir(os.path.dirname(opt.checkpoint))
torch.save(state, opt.checkpoint + '/%s.pth' % best_map_)
def main(**kwargs):
opt._parse(kwargs)
vis = Visualizer(env=opt.env)
# Model
print('==> Building model..')
net = DSOD(num_classes=21)
start_epoch = 0 # start from epoch 0 or last epoch
# Dataset
print('==> Preparing dataset..')
box_coder = SSDBoxCoder(net)
trainset = ListDataset(root=opt.data_root,
list_file=[opt.voc07_trainval, opt.voc12_trainval],
transform=Transform(box_coder, True))
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True)
net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
if opt.load_path is not None:
print('==> Resuming from checkpoint..')
checkpoint = torch.load(opt.load_path)
net.module.load_state_dict(checkpoint['net'])
criterion = SSDLoss(num_classes=21)
optimizer = optim.SGD(net.parameters(),
lr=opt.lr,
momentum=0.9,
weight_decay=5e-4)
best_map_ = 0
best_loss = 1e100
for epoch in range(start_epoch, start_epoch + 200):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
optimizer.zero_grad()
ix = 0
for batch_idx, (inputs, loc_targets,
cls_targets) in tqdm(enumerate(trainloader)):
inputs = Variable(inputs.cuda())
loc_targets = Variable(loc_targets.cuda())
cls_targets = Variable(cls_targets.cuda())
loc_preds, cls_preds = net(inputs)
ix += 1
loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets)
loss.backward()
train_loss += loss.data[0]
if (batch_idx + 1) % (opt.iter_size) == 0:
# if True:
for name, p in net.named_parameters():
p.grad.data.div_(ix)
ix = 0
optimizer.step()
optimizer.zero_grad()
if (batch_idx + 1) % opt.plot_every == 0:
vis.plot('loss', train_loss / (batch_idx + 1))
img = predict(
net, box_coder,
os.path.join(opt.data_root,
trainset.fnames[batch_idx]))
vis.img('predict', np.array(img).transpose(2, 0, 1))
if os.path.exists(opt.debug_file):
import ipdb
ipdb.set_trace()
# if (epoch+1)%10 == 0 :
# state = {
# 'net': net.module.state_dict(),
# # 'map': best_map_,
# 'epoch': epoch,
# }
# torch.save(state, opt.checkpoint + '/%s.pth' % epoch)
# if (epoch+1) % 30 == 0:
# for param_group in optimizer.param_groups:
# param_group['lr'] *= 0.1
current_loss = train_loss / (1 + batch_idx)
if current_loss < best_loss:
best_loss = current_loss
torch.save(net.module.state_dict(), '/tmp/dsod.pth')
if (epoch + 1) % opt.eval_every == 0:
net.module.load_state_dict(torch.load('/tmp/dsod.pth'))
aps = eval(net.module)
map_ = aps['map']
if map_ > best_map_:
print('Saving..')
state = {
'net': net.module.state_dict(),
'map': best_map_,
'epoch': epoch,
}
best_map_ = map_
if not os.path.isdir(os.path.dirname(opt.checkpoint)):
os.mkdir(os.path.dirname(opt.checkpoint))
best_path = opt.checkpoint + '/%s.pth' % best_map_
torch.save(state, best_path)
else:
net.module.load_state_dict(torch.load(best_path)['net'])
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.1
vis.log(
dict(epoch=(epoch + 1),
map=map_,
loss=train_loss / (batch_idx + 1)))
checkpoint = '/home/claude.cy/file/ssd/checkpoint/ckpt.pth'
data_root = '/home/claude.cy/.data/all_images'
voc07_trainval = 'torchcv/datasets/voc/voc07_trainval.txt'
voc12_trainval = 'torchcv/datasets/voc/voc12_trainval.txt'
voc07_test = 'torchcv/datasets/voc/voc07_test.txt'
batch_size = 16
num_worker = 8
args = Config()
import ipdb
ipdb.set_trace()
# Model
print('==> Building model..')
net = DSOD(num_classes=21)
start_epoch = 0 # start from epoch 0 or last epoch
if args.resume:
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.checkpoint)
net.load_state_dict(checkpoint['net'])
best_loss = checkpoint['loss']
start_epoch = checkpoint['epoch']
# Dataset
print('==> Preparing dataset..')
box_coder = SSDBoxCoder(net)
img_size = 300
def transform_train(img, boxes, labels):
class Detect():
def __init__(self, load_path):
self.img = None
print('Loading model..')
self.net = DSOD(num_classes=10)
self.net.cuda()
self.net = torch.nn.DataParallel(self.net,
device_ids=range(
torch.cuda.device_count()))
self.net.load_state_dict(torch.load(load_path)['net'], False)
self.net.eval()
self.box_coder = SSDBoxCoder(self.net.module)
def caffe_normalize(self, x):
return transforms.Compose([
transforms.Lambda(lambda x: 255 * x[[2, 1, 0]]),
transforms.Normalize([104, 117, 123],
(1, 1, 1)), # make it the same as caffe
# bgr and 0-255
])(x)
def transform(self, img, boxes):
img, boxes = resize(img, boxes, size=(opt.img_size, opt.img_size))
img = transforms.Compose([transforms.ToTensor(),
self.caffe_normalize])(img)
return img, boxes
def py_cpu_nms(self, dets, score, thresh):
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
scores = score
keep = []
index = scores.argsort()[::-1]
while index.size > 0:
i = index[
0] # every time the first is the biggst, and add it directly
keep.append(i)
# calculate the points of overlap
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1) # the weights of overlap
h = np.maximum(0, y22 - y11 + 1) # the height of overlap
overlaps = w * h
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= thresh)[0]
# print(idx)
index = index[idx + 1] # because index start from 1
return keep
def detect(self, img):
print('Processing img..')
self.img = Image.open(img)
if self.img.mode != 'RGB':
self.img = self.img.convert('RGB')
boxes = None
inputs, boxes = self.transform(self.img, boxes)
inputs = inputs.unsqueeze(0)
with torch.no_grad():
loc_preds, cls_preds = self.net(Variable(inputs.cuda()))
box_preds, label_preds, score_preds = self.box_coder.decode(
loc_preds.cpu().data.squeeze(),
F.softmax(cls_preds.squeeze(), dim=1).cpu().data,
score_thresh=0.1)
score = score_preds.numpy()
keep = self.py_cpu_nms(box_preds.numpy(), score, thresh=0.3)
box_preds = box_preds.numpy()[keep]
score_preds = score_preds.numpy()[keep]
label_preds = label_preds.numpy()[keep]
sw = float(self.img.size[0]) / float(opt.img_size)
sh = float(self.img.size[1]) / float(opt.img_size)
boxes = box_preds * np.array([sw, sh, sw, sh])
index = np.argmax(score_preds)
x1 = int(boxes[index][0])
x2 = int(boxes[index][2])
y1 = int(boxes[index][1])
y2 = int(boxes[index][3])
if x1 < 0: x1 = 0
if y1 < 0: y1 = 0
if x2 < 0: x2 = 0
if y2 < 0: y2 = 0
if x1 >= x2 or y1 >= y2:
print('no detection')
else:
return x1, y1, x2, y2