def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
def burnin_schedule(i):
if i < burn_in:
factor = pow(i / burn_in, 4)
elif i < steps[0]:
factor = 1.0
elif i < steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
if args.tfboard:
print("using tfboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(args.tfboard)
model.train()
imgsize = cfg['TRAIN']['IMGSIZE']
dataset = COCODataset(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params':value, 'weight_decay':decay * batch_size * subdivision}]
else:
params += [{'params':value, 'weight_decay':0.0}]
optimizer = optim.SGD(params, lr=base_lr, momentum=momentum,
dampening=0, weight_decay=decay * batch_size * subdivision)
iter_state = 0
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
iter_state = state['iter'] + 1
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
# start training loop
for iter_i in range(iter_state, iter_size + 1):
# COCO evaluation
if iter_i % args.eval_interval == 0 and iter_i > 0:
ap50_95, ap50 = evaluator.evaluate(model)
model.train()
if args.tfboard:
tblogger.add_scalar('val/COCOAP50', ap50, iter_i)
tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i)
# subdivision loop
optimizer.zero_grad()
for inner_iter_i in range(subdivision):
try:
imgs, targets, _, _ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, _, _ = next(dataiterator) # load a batch
imgs = Variable(imgs.type(dtype))
targets = Variable(targets.type(dtype), requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
scheduler.step()
if iter_i % 10 == 0:
# logging
current_lr = scheduler.get_lr()[0] * batch_size * subdivision
print('[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]'
% (iter_i, iter_size, current_lr,
model.loss_dict['xy'], model.loss_dict['wh'],
model.loss_dict['conf'], model.loss_dict['cls'],
model.loss_dict['l2'], imgsize),
flush=True)
if args.tfboard:
tblogger.add_scalar('train/total_loss', model.loss_dict['l2'], iter_i)
# random resizing
if random_resize:
imgsize = (random.randint(0, 9) % 10 + 10) * 32
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
# save checkpoint
if iter_i > 0 and (iter_i % args.checkpoint_interval == 0):
torch.save({'iter': iter_i,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
},
os.path.join(args.checkpoint_dir, "snapshot"+str(iter_i)+".ckpt"))
if args.tfboard:
tblogger.close()
def main():
"""
Visualize the detection result for the given image and the pre-trained model.
"""
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=0)
parser.add_argument('--cfg', type=str, default='config/yolov3_default.cfg')
parser.add_argument('--ckpt', type=str, help='path to the checkpoint file')
parser.add_argument('--weights_path',
type=str,
default=None,
help='path to weights file')
parser.add_argument('--image', type=str)
parser.add_argument(
'--background',
action='store_true',
default=False,
help='background(no-display mode. save "./output.png")')
parser.add_argument('--detect_thresh',
type=float,
default=None,
help='confidence threshold')
args = parser.parse_args()
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
imgsize = cfg['TEST']['IMGSIZE']
model = YOLOv3(cfg['MODEL'])
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
if args.detect_thresh:
confthre = args.detect_thresh
img = cv2.imread(args.image)
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
img, info_img = preprocess(img, imgsize,
jitter=0) # info = (h, w, nh, nw, dx, dy)
img = np.transpose(img / 255., (2, 0, 1))
img = torch.from_numpy(img).float().unsqueeze(0)
if args.gpu >= 0:
model.cuda(args.gpu)
img = Variable(img.type(torch.cuda.FloatTensor))
else:
img = Variable(img.type(torch.FloatTensor))
if args.weights_path:
print("loading yolo weights %s" % (args.weights_path))
parse_yolo_weights(model, args.weights_path)
else:
print("loading checkpoint %s" % (args.ckpt))
model.load_state_dict(torch.load(args.ckpt))
model.eval()
with torch.no_grad():
outputs = model(img)
outputs = postprocess(outputs, 80, confthre, nmsthre)
if outputs[0] is None:
print("No Objects Deteted!!")
return
coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names(
)
bboxes = list()
classes = list()
colors = list()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = coco_class_ids[int(cls_pred)]
print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
print('\t+ Label: %s, Conf: %.5f' %
(coco_class_names[cls_id], cls_conf.item()))
box = yolobox2label([y1, x1, y2, x2], info_img)
bboxes.append(box)
classes.append(cls_id)
colors.append(coco_class_colors[int(cls_pred)])
if args.background:
import matplotlib
matplotlib.use('Agg')
from utils.vis_bbox import vis_bbox
import matplotlib.pyplot as plt
vis_bbox(img_raw,
bboxes,
label=classes,
label_names=coco_class_names,
instance_colors=colors,
linewidth=2)
#plt.show()
if args.background:
plt.savefig('output.png')
def main():
args = parse_args()
print("------------------------------------")
print(" use {} dataset for demo. ".format(args.data))
print("------------------------------------")
assert args.data in ['coco', 'drone']
if torch.cuda.is_available() and args.gpu >= 0:
device = torch.device('cuda:{}'.format(args.gpu))
else:
device = torch.device('cpu')
# [TBM] gen n_classes from coco-format json file..
if args.data == 'coco':
cfg_path = 'config/yolov3_default.cfg'
n_classes = 80
if args.data == 'drone':
cfg_path = 'config/yolov3_visdrone_default.cfg'
n_classes = 10
with open(cfg_path, 'r') as f:
cfg = yaml.load(f)
imgsize = cfg['TEST']['IMGSIZE']
model = YOLOv3(n_classes=n_classes)
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
if args.detect_thresh:
confthre = args.detect_thresh
img = cv2.imread(args.image)
assert img is not None
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
img, info_img = preprocess(img, imgsize) # info = (h, w, nh, nw, dx, dy)
img = torch.from_numpy(img).float().unsqueeze(0)
model = model.to(device)
img = Variable(img.to(device, dtype=torch.float32))
if args.weights_path:
print("loading yolo weights %s" % (args.weights_path))
parse_yolo_weights(model, args.weights_path)
else:
print("loading checkpoint %s" % (args.ckpt))
model.load_state_dict(torch.load(args.ckpt))
model.eval()
with torch.no_grad():
outputs = model(img)
outputs = postprocess(outputs, n_classes, confthre, nmsthre)
# [TBM] gen label_names from coco-format json file..
if args.data == 'coco':
class_names, class_ids, class_colors = get_coco_label_names()
if args.data == 'drone':
class_names, class_ids, class_colors = get_visdrone_label_names()
bboxes = list()
classes = list()
colors = list()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = class_ids[int(cls_pred)]
print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
print('\t+ Label: %s, Conf: %.5f' %
(class_names[cls_id], cls_conf.item()))
box = yolobox2label([y1, x1, y2, x2], info_img)
bboxes.append(box)
classes.append(cls_id)
colors.append(class_colors[int(cls_pred)])
vis_bbox(img_raw,
bboxes,
label=classes,
label_names=class_names,
instance_colors=colors,
linewidth=2)
if args.window:
plt.show()
else:
out_path = './output.png'
plt.savefig(out_path, bbox_inches=0, pad_inches=0, dpi=100)
def main(
image = None ,
gpu = -1,
weights_path= f"{ Path(__file__).parent }/weights/yolov3.weights",
background = False
):
"""
Visualize the detection result for the given image and the pre-trained model.
"""
print( weights_path )
my_path = Path( __file__ ).parent
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default= gpu )
parser.add_argument('--cfg', type=str, default=my_path/'config/yolov3_default.cfg')
parser.add_argument('--ckpt', type=str,
help='path to the checkpoint file')
parser.add_argument('--weights_path', type=str,
default= weights_path, help='path to weights file')
parser.add_argument('--image', type=str , default= image )
parser.add_argument('--background', type=bool,
default= background , help='background(no-display mode. save "./output.png")')
parser.add_argument('--detect_thresh', type=float,
default= 0.5 , help='confidence threshold')
args = parser.parse_args()
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
imgsize = cfg['TEST']['IMGSIZE']
model = YOLOv3(cfg['MODEL'])
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
if args.detect_thresh:
confthre = args.detect_thresh
img = imread( args.image )
if img is None :
print( "load image failed" )
print( args.image )
return
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
img, info_img = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy)
img = np.transpose(img / 255., (2, 0, 1))
img = torch.from_numpy(img).float().unsqueeze(0)
if args.gpu >= 0:
model.cuda(args.gpu)
img = Variable(img.type(torch.cuda.FloatTensor))
else:
img = Variable(img.type(torch.FloatTensor))
assert args.weights_path or args.ckpt, 'One of --weights_path and --ckpt must be specified'
if args.weights_path:
print("loading yolo weights %s" % (args.weights_path))
parse_yolo_weights(model, args.weights_path)
elif args.ckpt:
print("loading checkpoint %s" % (args.ckpt))
state = torch.load(args.ckpt)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
model.eval()
with torch.no_grad():
outputs1 = model(img)
# np.save("output.npy" , outputs.numpy() )
# torch.save( outputs1 , "outputs1.pt" )
out1 = torch.load( "outputs1.pt" )
rere = torch.equal( outputs1 , out1 )
outputs = postprocess(outputs1, 80, confthre, nmsthre)
a = "hoho"
if outputs[0] is None:
print("No Objects Deteted!!")
return
coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names()
bboxes = list()
classes = list()
colors = list()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = coco_class_ids[int(cls_pred)]
print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
print('\t+ Label: %s, Conf: %.5f' %
(coco_class_names[cls_id], cls_conf.item()))
box = yolobox2label([y1, x1, y2, x2], info_img)
bboxes.append(box)
classes.append(cls_id)
colors.append(coco_class_colors[int(cls_pred)])
# args.background = True
if args.background:
import matplotlib
matplotlib.use('Agg')
from utils.vis_bbox import vis_bbox
vis_bbox(
img_raw, bboxes, label=classes, label_names=coco_class_names,
instance_colors=colors, linewidth=2)
if args.background:
output = Path( "./output" )
output.mkdir( parents=True , exist_ok=True )
now = datetime.now().strftime("%Y-%m-%d %H-%M-%S")
output /= f"output-{now}.png"
plt.savefig( output )
return str( output.absolute() )
# return plt_to_qpixmap(plt.gca())
else :
plt.show()
from utils.utils import *
from utils.parse_yolo_weights import parse_yolo_weights
from person_det import detect
from mask_function import mask_plot
with open('config/yolov3_default.cfg', 'r') as f:
cfg = yaml.load(f)
imgsize_y = cfg['TEST']['IMGSIZE']
model_Y = YOLOv3(cfg['MODEL'])
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
model_Y.cuda(0)
parse_yolo_weights(model_Y, 'weights/yolov3.weights')
model_Y.eval()
####################################################
mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3)
std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3)
mpii_edges = [[0, 1], [0, 2], [1, 3], [2, 4], [4, 6], [3, 5], [5, 6], [5, 7], [7, 9], [6, 8], [8, 10], [6, 12], [5, 11], [11, 12], [12, 14], [14, 16], [11, 13], [13, 15]]
def show_2d(pts, c, edges,orig_image,min_x,min_y):
for i in range(len(pts)):
points = pts[i]
num_joints = points.shape[0]
points = ((points.reshape(num_joints, -1))).astype(np.int32)
#for j in range(num_joints):
# cv2.circle(img, (points[j, 0], points[j, 1]), 3, c, -1)
for e in edges:
def main():
"""
Visualize the detection result for the given image and the pre-trained model.
"""
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=0)
args = parser.parse_args()
with open('config/yolov3_default.cfg', 'r') as f:
cfg = yaml.load(f)
imgsize = cfg['TEST']['IMGSIZE']
model = YOLOv3(cfg['MODEL'])
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
img = cv2.imread('data/mpii_87.png')
orig_img = img.copy()
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
img, info_img = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy)
img = np.transpose(img / 255., (2, 0, 1))
img = torch.from_numpy(img).float().unsqueeze(0)
if args.gpu >= 0:
model.cuda(args.gpu)
img = Variable(img.type(torch.cuda.FloatTensor))
else:
img = Variable(img.type(torch.FloatTensor))
parse_yolo_weights(model, 'weights/yolov3.weights')
model.eval()
with torch.no_grad():
outputs = model(img)
outputs = postprocess(outputs, 80, confthre, nmsthre)
if outputs[0] is None:
print("No Objects Deteted!!")
return
coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names()
bboxes = list()
classes = list()
colors = list()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = coco_class_ids[int(cls_pred)]
if cls_id==1:
print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
print('\t+ Label: %s, Conf: %.5f' %
(coco_class_names[cls_id], cls_conf.item()))
box = yolobox2label([y1, x1, y2, x2], info_img)
bboxes.append(box)
y_min = int(box[0])
x_min = int(box[1])
y_max = int(box[2])
x_max = int(box[3])
#cropped = orig_img[y_min:y_max,x_min:x_max]
cv2.rectangle(orig_img,(x_min,y_min),(x_max,y_max),(0,0,255),2)
cv2.imshow('img',orig_img)
cv2.waitKey(0)
def main():
"""
Visualize the detection result for the given image and the pre-trained model.
"""
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=0)
parser.add_argument('--cfg',
type=str,
default='config/yolov3_default_digestpath.cfg')
parser.add_argument('--ckpt', type=str, help='path to the checkpoint file')
parser.add_argument('--weights_path',
type=str,
default=None,
help='path to weights file')
parser.add_argument('--image', type=str)
parser.add_argument(
'--background',
action='store_true',
default=False,
help='background(no-display mode. save "./output.png")')
parser.add_argument('--detect_thresh',
type=float,
default=None,
help='confidence threshold')
parser.add_argument('--dataset',
help='dataset to work with: {}'.format(
Dataset.print_choices()),
type=int,
default=Dataset.SIGNET_RING)
args = parser.parse_args()
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
imgsize = cfg['TEST']['IMGSIZE']
model = YOLOv3(cfg['MODEL'])
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
if args.detect_thresh:
confthre = args.detect_thresh
img = cv2.imread(args.image)
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
img, info_img = preprocess(img, imgsize,
jitter=0) # info = (h, w, nh, nw, dx, dy)
img = np.transpose(img / 255., (2, 0, 1))
img = torch.from_numpy(img).float().unsqueeze(0)
if args.gpu >= 0:
model.cuda(args.gpu)
img = Variable(img.type(torch.cuda.FloatTensor))
else:
img = Variable(img.type(torch.FloatTensor))
assert args.weights_path or args.ckpt, 'One of --weights_path and --ckpt must be specified'
if args.weights_path:
print("loading yolo weights %s" % (args.weights_path))
parse_yolo_weights(model, args.weights_path)
elif args.ckpt:
print("loading checkpoint %s" % (args.ckpt))
state = torch.load(args.ckpt)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
model.eval()
with torch.no_grad():
outputs = model(img)
outputs = postprocess(outputs, Dataset.NUM_CLASSES[args.dataset],
confthre, nmsthre)
if outputs[0] is None:
print("No Objects Deteted!!")
return
bboxes = list()
colors = list()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
print('\t+ Conf: %.5f' % cls_conf.item())
box = yolobox2label([y1, x1, y2, x2], info_img)
bboxes.append(box)
colors.append(BOX_COLOR)
if args.background:
matplotlib.use('Agg')
vis_bbox(img_raw, bboxes, instance_colors=colors, linewidth=2)
plt.show()
if args.background:
plt.savefig('output.png')
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision
datatype = cfg['TRAIN']['DATATYPE']
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
def burnin_schedule(i):
if i < burn_in:
factor = pow(i / burn_in, 4)
elif i < steps[0]:
factor = 1.0
elif i < steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
if args.tfboard:
print("using tfboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(args.tfboard)
model.train()
coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names()
imgsize = cfg['TRAIN']['IMGSIZE']
if datatype=='voc':
dataset = VOCDataset(model_type=cfg['MODEL']['TYPE'],
data_dir='../../VOCdevkit/VOC2007',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
print('load voc dataset successfully')
else:
dataset = COCODataset(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
print('load COCO dataset successfully')
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params': value, 'weight_decay': decay
* batch_size * subdivision}]
else:
params += [{'params': value, 'weight_decay': 0.0}]
optimizer = optim.SGD(params, lr=base_lr, momentum=momentum,
dampening=0, weight_decay=decay * batch_size * subdivision)
iter_state = 0
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
iter_state = state['iter'] + 1
#学习率控制 Sets the learning rate of each parameter group to the initial lr times a given function. When last_epoch=-1, sets initial lr as lr.
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
# result=evals(model)
# print(result)
# start training loop
# print('args.eval_interval',args.eval_interval)
for iter_i in range(iter_state, iter_size + 1):
if iter_i % (args.eval_interval*2) == 0 and iter_i > 0:
if datatype=='voc':
result=evals(model)
print(result)
else:
ap50_95, ap50 = evaluator.evaluate(model)
print(ap50_95, ap50)
model.train()
if args.tfboard:
tblogger.add_scalar('val/COCOAP50', ap50, iter_i)
tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i)
if iter_i % (40000) == 0 and iter_i > 0:
draw(model,datatype,imgsize)
# subdivision loop
optimizer.zero_grad()
for inner_iter_i in range(subdivision):
try:
imgs, targets, info_img, id_ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, info_img, id_ = next(dataiterator) # load a batch
imgs = Variable(imgs.type(dtype))
targets = Variable(targets.type(dtype), requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
scheduler.step()
if iter_i % 10 == 0:
# logging
current_lr = scheduler.get_lr()[0] * batch_size * subdivision
print('[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]'
% (iter_i, iter_size, current_lr,
model.loss_dict['xy'], model.loss_dict['wh'],
model.loss_dict['conf'], model.loss_dict['cls'],
model.loss_dict['l2'], imgsize))
if args.tfboard:
tblogger.add_scalar('train/total_loss',
model.loss_dict['l2'], iter_i)
# random resizing
# 变输入大小,利用了yolov3网络的全卷积，使得模型不受图像大小的改变而影响参数。
if random_resize:
imgsize = (random.randint(0, 9) % 10 + 10) * 32
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
if iter_i % 100 == 0:
model.eval()
if datatype=='voc':
img_file = os.path.join('../../VOCdevkit/VOC2007', 'JPEGImages', id_[0] + '.jpg')
else:
img_file = os.path.join('COCO', 'train2017',
'{:012}'.format(id_) + '.jpg')
img = cv2.imread(img_file)
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
# print(img_raw.shape)
# print(img_raw)
# print(imgs)
img, info_img_t = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy)
img = np.transpose(img / 255., (2, 0, 1))
img = torch.from_numpy(img).float().unsqueeze(0)
img = Variable(img.type(torch.cuda.FloatTensor))
outputs = model(img)
#outputs.shape : torch.Size([1, 12348, 85])
outputs = postprocess(outputs, 80, 0.5, 0.5)
# imgs.shape : torch.Size([1, 3, 608, 608])
# outputs[0].shape :torch.Size([3, 7])
# targets.shape :torch.Size([1, 50, 5])
# print(outputs)
if outputs[0] is not None:
bboxes = list()
classes = list()
colors = list()
# print(info_img_t)
info_img=tuple(info_img)
# print(info_img)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = coco_class_ids[int(cls_pred)]
# print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
# print('\t+ Label: %s, Conf: %.5f' %
# (coco_class_names[cls_id], cls_conf.item()))
# print([y1, x1, y2, x2])
box = yolobox2label([y1, x1, y2, x2], info_img_t)
bboxes.append(box)
classes.append(cls_id)
colors.append(coco_class_colors[int(cls_pred)])
vis_bbox(
img_raw, bboxes, label=classes, label_names=coco_class_names,
instance_colors=colors, linewidth=2)
plt.savefig('output/'+str(iter_i)+'.jpg')
model.train()
# save checkpoint
if iter_i > 0 and (iter_i % args.checkpoint_interval == 0):
torch.save({'iter': iter_i,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
},
os.path.join(args.checkpoint_dir, "snapshot" + str(iter_i) + ".ckpt"))
if args.tfboard:
tblogger.close()
def main():
args = parse_args()
print("------------------------------------")
print(" use {} dataset for demo. ".format(args.data))
print("------------------------------------")
assert args.data in ['coco', 'drone']
if torch.cuda.is_available() and args.gpu >= 0:
device = torch.device('cuda:{}'.format(args.gpu))
else:
device = torch.device('cpu')
# [TBM] gen n_classes from coco-format json file..
if args.data == 'coco':
cfg_path = 'config/yolov3_default.cfg'
n_classes = 80
if args.data == 'drone':
cfg_path = 'config/yolov3_visdrone_default.cfg'
n_classes = 10
with open(cfg_path, 'r') as f:
cfg = yaml.load(f)
imgsize = cfg['TEST']['IMGSIZE']
model = YOLOv3(n_classes=n_classes)
confthre = cfg['TEST']['CONFTHRE']
nmsthre = cfg['TEST']['NMSTHRE']
if args.detect_thresh:
confthre = args.detect_thresh
model = model.to(device)
if args.weights_path:
print("loading yolo weights %s" % (args.weights_path))
parse_yolo_weights(model, args.weights_path)
else:
print("loading checkpoint %s" % (args.ckpt))
model.load_state_dict(torch.load(args.ckpt))
model.eval()
dir_name = os.path.basename(os.path.dirname(args.in_dir + '/'))
out_dir = os.path.join(args.out_dir, dir_name)
os.makedirs(out_dir, exist_ok=True)
img_files = os.listdir(args.in_dir)
img_files.sort()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
for i in range(0, len(img_files), args.step):
filename = img_files[i]
img_path = os.path.join(args.in_dir, filename)
img = cv2.imread(img_path)
assert img is not None
start.record()
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
img, info_img = preprocess(img,
imgsize) # info = (h, w, nh, nw, dx, dy)
img = torch.from_numpy(img).float().unsqueeze(0)
img = Variable(img.to(device, dtype=torch.float32))
with torch.no_grad():
outputs = model(img)
outputs = postprocess(outputs, n_classes, confthre, nmsthre)
end.record()
torch.cuda.synchronize()
# [TBM] gen label_names from coco-format json file..
if args.data == 'coco':
class_names, class_ids, class_colors = get_coco_label_names()
if args.data == 'drone':
class_names, class_ids, class_colors = get_visdrone_label_names()
bboxes, classes, colors = list(), list(), list()
if outputs[0] is None:
outputs[0] = []
if args.verbose:
print("=====================================")
print("{}, {:.2f} [fps]".format(filename,
1000.0 / start.elapsed_time(end)))
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = class_ids[int(cls_pred)]
if args.verbose:
print(int(x1), int(y1), int(x2), int(y2), float(conf),
int(cls_pred))
print('\t+ Label: %s, Conf: %.5f' %
(class_names[cls_id], cls_conf.item()))
box = yolobox2label([y1, x1, y2, x2], info_img)
bboxes.append(box)
classes.append(cls_id)
colors.append(class_colors[int(cls_pred)])
if args.verbose:
print()
vis_bbox(img_raw,
bboxes,
label=classes,
label_names=class_names,
instance_colors=colors,
linewidth=2)
basename, _ = os.path.splitext(filename)
out_path = os.path.join(out_dir, '{}.png'.format(basename))
plt.savefig(out_path, bbox_inches=0, pad_inches=0, dpi=100)
plt.close()
print("Done!")
