def get_model(opt, pretrained=None, trn=True, weights_FIMs=None, alpha=1.):
'''Getting model and initializing.
Args:
pretrained, None or path to pretrained model weights
trn, True for training and False for evaluating'''
# Model structure
model = Darknet(opt.model_config_path, opt.img_size, weights_FIMs, alpha)
print(model)
# Initialize
model.apply(weights_init_normal)
# Pretrained or not
coco_weights = True if pretrained == 'weights/yolov3.weights' else False
try:
model.load_weights(pretrained, use_coco=coco_weights)
except TypeError:
pass
# Cuda or not
if opt.cuda:
model = model.cuda()
cudnn.benchmark = True
# Mode = train or eval
if trn:
model.train()
else:
model.eval()
return model
def eval_flowchart(init_style, para_part, reg, alpha, ablation_type):
'''Main body for evaluation'''
args = all_args()
# Storage path 'eval/'
os.makedirs(args.save_folder, exist_ok=True)
# Dataset
dataset = get_data(args)
# Load net
net = Darknet(args.model_config_path, img_size=args.img_size)
# Visdom
viz, epoch_aps = init_viz(args, init_style, para_part, reg, alpha, dataset)
# Evaluate
ckpt_path, ckpts = get_ckpts(args, init_style, para_part, reg, alpha,
ablation_type)
mAP_max = 0
for ckpt_idx, ckpt in enumerate(ckpts):
# sample one for hyperparameter adjustment
if ckpt_idx < 120:
continue
# Make output dir
dir_name = '_'.join([
ablation_type, args.arc, args.dataset, args.set_type, init_style,
para_part, reg,
str(alpha), ckpt,
str(ckpt_idx)
])
output_dir = get_output_dir(args.save_folder, dir_name)
# Load weight
args.weight_path = os.path.join(ckpt_path, ckpt)
# assert os.path.isfile(args.weight_path)
try:
net.load_weights(args.weight_path)
except FileNotFoundError as err:
print(err)
# Cuda or not
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
net.eval()
print('Finished loading model!')
# Evaluation, use_07_eval False
aps, mAP = test_net(output_dir,
net,
args.cuda,
dataset,
args.score_thres,
args.nms_thres,
use_07_eval=False,
iou_thres=args.iou_thres)
# If not greater than before, delete
if mAP_max >= mAP:
rmtree(output_dir)
else:
mAP_max = mAP
# Visdom
update_vis(viz, epoch_aps, ckpt_idx + 1, *aps, mAP)
def main():
# model
model = Darknet(str(args.config), img_size=416)
model_wts = torch.load(str(args.weight), map_location='cpu')
model.load_state_dict(model_wts)
if torch.cuda.is_available():
print('gpu: available')
model = model.cuda()
else:
print('gpu: not available')
# read video
print(">> reading video...")
video, info = read_video(args.input)
video = video[:, :, :, ::-1]
print(" shape:", video.shape)
print(" info: ", info)
# forward
print(">> predicting...")
imgs, bboxes, ss, labels = [], [], [], []
for i in tqdm(range(0, len(video))):
img = video[i]
bbox, max_s = model.predict(img, args.conf_thresh, args.nms_thresh)
imgs.append(img)
ss.append(max_s)
if len(bbox) != 0:
bboxes.append([bbox])
labels.append([0])
else:
bboxes.append([])
labels.append([])
# draw bbox
imgs = np.asarray(imgs)
# imgs = imgs[:,:,::-1]
for i in tqdm(range(len(imgs))):
if len(bboxes[i]) != 0:
ty, tx, by, bx = [int(n) for n in bboxes[i][0]]
cv2.rectangle(imgs[i], (tx, ty), (bx, by), (255, 0, 0), 3)
# save as video
print(">> saving video...")
imgs = imgs[:, :, :, ::-1]
write_video(args.output, imgs, info)
learning_rate = float(hyperparams['learning_rate'])
momentum = float(hyperparams['momentum'])
decay = float(hyperparams['decay'])
burn_in = int(hyperparams['burn_in'])
# Initiate model
dirlist = os.listdir('./yolodata/')
dirlist.remove('showTruth.m')
for dir in sorted(dirlist):
# Get dataloader
model = Darknet(opt.model_config_path)
model.load_weights(opt.weights_path)
if cuda:
model = model.cuda()
model.train()
train_path = './yolodata/' + dir + '/train/indextrain.txt'
dataloader = torch.utils.data.DataLoader(ListDataset(train_path),
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.n_cpu)
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
optimizer = optim.SGD(model.parameters(),
lr=learning_rate,
momentum=momentum,
dampening=0,
weight_decay=decay)
def train(
cfg,
data_cfg,
weights_from="",
weights_to="",
save_every=10,
img_size=(1088, 608),
resume=False,
epochs=100,
batch_size=16,
accumulated_batches=1,
freeze_backbone=False,
opt=None,
):
# The function starts
NUM_WORKERS = opt.num_workers
timme = strftime("%Y-%d-%m %H:%M:%S", gmtime())
timme = timme[5:-3].replace('-', '_')
timme = timme.replace(' ', '_')
timme = timme.replace(':', '_')
weights_to = osp.join(weights_to, 'run' + timme)
mkdir_if_missing(weights_to)
mkdir_if_missing(weights_to + '/cfg/')
if resume:
latest_resume = osp.join(weights_from, 'latest.pt')
torch.backends.cudnn.benchmark = True # unsuitable for multiscale
# Configure run
f = open(data_cfg)
data_config = json.load(f)
trainset_paths = data_config['train']
dataset_root = data_config['root']
f.close()
transforms = T.Compose([T.ToTensor()])
# Get dataloader
dataset = JointDataset(dataset_root,
trainset_paths,
img_size,
augment=True,
transforms=transforms)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=NUM_WORKERS,
pin_memory=True,
drop_last=True,
collate_fn=collate_fn)
# Initialize model
model = Darknet(cfg, dataset.nID)
cutoff = -1 # backbone reaches to cutoff layer
start_epoch = 0
if resume:
checkpoint = torch.load(latest_resume, map_location='cpu')
# Load weights to resume from
model.load_state_dict(checkpoint['model'])
model.cuda().train()
# Set optimizer
optimizer = torch.optim.SGD(filter(lambda x: x.requires_grad,
model.parameters()),
lr=opt.lr,
momentum=.9)
start_epoch = checkpoint['epoch'] + 1
if checkpoint['optimizer'] is not None:
optimizer.load_state_dict(checkpoint['optimizer'])
del checkpoint # current, saved
else:
# Initialize model with backbone (optional)
if cfg.endswith('yolov3.cfg'):
load_darknet_weights(model,
osp.join(weights_from, 'darknet53.conv.74'))
cutoff = 75
elif cfg.endswith('yolov3-tiny.cfg'):
load_darknet_weights(model,
osp.join(weights_from, 'yolov3-tiny.conv.15'))
cutoff = 15
model.cuda().train()
# Set optimizer
optimizer = torch.optim.SGD(filter(lambda x: x.requires_grad,
model.parameters()),
lr=opt.lr,
momentum=.9,
weight_decay=1e-4)
model = torch.nn.DataParallel(model)
# Set scheduler
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(0.5 * opt.epochs),
int(0.75 * opt.epochs)],
gamma=0.1)
# An important trick for detection: freeze bn during fine-tuning
if not opt.unfreeze_bn:
for i, (name, p) in enumerate(model.named_parameters()):
p.requires_grad = False if 'batch_norm' in name else True
# model_info(model)
t0 = time.time()
for epoch in range(epochs):
epoch += start_epoch
logger.info(
('%8s%12s' + '%10s' * 6) % ('Epoch', 'Batch', 'box', 'conf', 'id',
'total', 'nTargets', 'time'))
# Freeze darknet53.conv.74 for first epoch
if freeze_backbone and (epoch < 2):
for i, (name, p) in enumerate(model.named_parameters()):
if int(name.split('.')[2]) < cutoff: # if layer < 75
p.requires_grad = False if (epoch == 0) else True
ui = -1
rloss = defaultdict(float) # running loss
## training schedule
optimizer.zero_grad()
for i, (imgs, targets, _, _, targets_len) in enumerate(dataloader):
if sum([len(x) for x in targets]) < 1: # if no targets continue
continue
# SGD burn-in
burnin = min(1000, len(dataloader))
if (epoch == 0) & (i <= burnin):
lr = opt.lr * (i / burnin)**4
for g in optimizer.param_groups:
g['lr'] = lr
# Compute loss, compute gradient, update parameters
loss, components = model(imgs.cuda(), targets.cuda(),
targets_len.cuda())
components = torch.mean(components.view(-1, 5), dim=0)
loss = torch.mean(loss)
loss.backward()
# accumulate gradient for x batches before optimizing
if ((i + 1) % accumulated_batches
== 0) or (i == len(dataloader) - 1):
optimizer.step()
optimizer.zero_grad()
# Running epoch-means of tracked metrics
ui += 1
for ii, key in enumerate(model.module.loss_names):
rloss[key] = (rloss[key] * ui + components[ii]) / (ui + 1)
# rloss indicates running loss values with mean updated at every epoch
s = ('%8s%12s' + '%10.3g' * 6) % (
'%g/%g' % (epoch, epochs - 1), '%g/%g' %
(i, len(dataloader) - 1), rloss['box'], rloss['conf'],
rloss['id'], rloss['loss'], rloss['nT'], time.time() - t0)
t0 = time.time()
if i % opt.print_interval == 0:
logger.info(s)
# Save latest checkpoint
checkpoint = {
'epoch': epoch,
'model': model.module.state_dict(),
'optimizer': optimizer.state_dict()
}
copyfile(cfg, weights_to + '/cfg/yolo3.cfg')
copyfile(data_cfg, weights_to + '/cfg/ccmcpe.json')
latest = osp.join(weights_to, 'latest.pt')
torch.save(checkpoint, latest)
if epoch % save_every == 0 and epoch != 0:
# making the checkpoint lite
checkpoint["optimizer"] = []
torch.save(
checkpoint,
osp.join(weights_to, "weights_epoch_" + str(epoch) + ".pt"))
# Calculate mAP
'''
if epoch % opt.test_interval == 0:
with torch.no_grad():
mAP, R, P = test.test(cfg, data_cfg, weights=latest, batch_size=batch_size, img_size=img_size,
print_interval=40, nID=dataset.nID)
test.test_emb(cfg, data_cfg, weights=latest, batch_size=batch_size, img_size=img_size,
print_interval=40, nID=dataset.nID)
'''
# Call scheduler.step() after opimizer.step() with pytorch > 1.1.0
scheduler.step()
class YOLOV4(object):
if CWD == THIS_DIR:
_defaults = {
"weights": "weights/yolov4.weights",
"config": "cfg/yolov4.cfg",
"classes_path": 'cfg/coco.names',
"thresh": 0.5,
"nms_thresh": 0.4,
"model_image_size": (608, 608),
"max_batch_size": 4,
"half": True
}
else:
_defaults = {
"weights": "yolov4_pytorch/weights/yolov4.weights",
"config": "yolov4_pytorch/cfg/yolov4.cfg",
"classes_path": 'yolov4_pytorch/cfg/coco.names',
"thresh": 0.5,
"nms_thresh": 0.4,
"model_image_size": (608, 608),
"max_batch_size": 4,
"half": True
}
def __init__(self, bgr=True, gpu_device=0, **kwargs):
self.__dict__.update(self._defaults) # set up default values
# for portability between keras-yolo3/yolo.py and this
if 'model_path' in kwargs:
kwargs['weights'] = kwargs['model_path']
if 'score' in kwargs:
kwargs['thresh'] = kwargs['score']
self.__dict__.update(kwargs) # update with user overrides
self.class_names = self._get_class()
self.model = Darknet(self.config)
self.model.load_darknet_weights(self.weights)
self.device = gpu_device
self.model.cuda(self.device)
self.model.eval()
self.bgr = bgr
if self.half:
self.model.half()
# warm up
self._detect([np.zeros((10, 10, 3), dtype=np.uint8)])
print('Warmed up!')
def _get_class(self):
with open(self.classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
def _detect(self, list_of_imgs):
inputs = []
for img in list_of_imgs:
if self.bgr:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# print('bgr: {}'.format(img.shape))
# print('size: {}'.format(self.model_image_size))
image = cv2.resize(img, self.model_image_size)
# print('image: {}'.format(image.shape))
inputs.append(np.expand_dims(np.array(image), axis=0))
images = np.concatenate(inputs, 0)
# print('images: {}'.format(images.shape))
images = torch.from_numpy(images.transpose(0, 3, 1,
2)).float().div(255.0)
images = images.cuda()
images = torch.autograd.Variable(images)
if self.half:
images = images.half()
batches = []
for i in range(0, len(images), self.max_batch_size):
these_imgs = images[i:i + self.max_batch_size]
batches.append(these_imgs)
feature_list = None
with torch.no_grad():
for batch in batches:
img = batch.cuda(self.device)
features = self.model(img)
if feature_list is None:
feature_list = features
else:
feature_list = torch.cat((feature_list, features))
# feature_list: (batch, height * width * num_anchors, 5 + num_classes)
return feature_list
def detect_get_box_in(self,
images,
box_format='ltrb',
classes=None,
buffer_ratio=0.0):
'''
Params
------
- images : ndarray-like or list of ndarray-like
- box_format : string of characters representing format order, where l = left, t = top, r = right, b = bottom, w = width and h = height
- classes : list of string, classes to focus on
- buffer : float, proportion of buffer around the width and height of the bounding box
Returns
-------
if one ndarray given, this returns a list (boxes in one image) of tuple (box_infos, score, predicted_class),
else if a list of ndarray given, this return a list (batch) containing the former as the elements,
where,
- box_infos : list of floats in the given box format
- score : float, confidence level of prediction
- predicted_class : string
'''
single = False
if isinstance(images, list):
if len(images) <= 0:
return None
else:
assert all(isinstance(im, np.ndarray) for im in images)
elif isinstance(images, np.ndarray):
images = [images]
single = True
res = self._detect(images)
frame_shapes = [image.shape for image in images]
all_dets = self._postprocess(res,
shapes=frame_shapes,
box_format=box_format,
classes=classes,
buffer_ratio=buffer_ratio)
if single:
return all_dets[0]
else:
return all_dets
def get_detections_dict(self, frames, classes=None, buffer_ratio=0.0):
'''
Params: frames, list of ndarray-like
Returns: detections, list of dict, whose key: label, confidence, t, l, w, h
'''
if frames is None or len(frames) == 0:
return None
all_dets = self.detect_get_box_in(frames,
box_format='tlbrwh',
classes=classes,
buffer_ratio=buffer_ratio)
all_detections = []
for dets in all_dets:
detections = []
for tlbrwh, confidence, label in dets:
top, left, bot, right, width, height = tlbrwh
detections.append({
'label': label,
'confidence': confidence,
't': top,
'l': left,
'b': bot,
'r': right,
'w': width,
'h': height
})
all_detections.append(detections)
return all_detections
def _nms(self, predictions):
predictions[..., :4] = self.xywh2p1p2(predictions[..., :4])
outputs = [None for _ in range(len(predictions))]
for i, image_pred in enumerate(predictions):
image_pred = image_pred[image_pred[:, 4] >= self.thresh]
# If none anchor are remaining => process next image
if not image_pred.size(0):
continue
# Object confidence times class confidence (n, ) * (n, )
score = image_pred[:, 4] * image_pred[:, 5:].max(1)[0]
class_confs, class_preds = image_pred[:, 5:].max(1, keepdim=True)
detections = torch.cat(
(image_pred[:, :5], class_confs.type(
predictions.dtype), class_preds.type(predictions.dtype)),
dim=1)
keep = batched_nms(image_pred[:, :4].float(), score,
class_preds[:, 0], self.nms_thresh)
outputs[i] = detections[keep]
return outputs
@staticmethod
def xywh2p1p2(x):
y = x.new(x.shape)
y[..., 0] = x[..., 0] - x[..., 2] / 2.
y[..., 1] = x[..., 1] - x[..., 3] / 2.
y[..., 2] = x[..., 0] + x[..., 2] / 2.
y[..., 3] = x[..., 1] + x[..., 3] / 2.
return y
@staticmethod
def p1p2Toxywh(x):
y = x.new(x.shape)
y[..., 0] = x[..., 0]
y[..., 1] = x[..., 1]
y[..., 2] = x[..., 2] - x[..., 0]
y[..., 3] = x[..., 3] - x[..., 1]
return y
def _postprocess(self,
outputs,
shapes,
box_format='ltrb',
classes=None,
buffer_ratio=0.0):
outputs = self._nms(outputs)
detections = []
for i, frame_bbs in enumerate(outputs):
im_height, im_width, _ = shapes[i]
if frame_bbs is None:
detections.append([])
continue
frame_bbs = self._resize_boxes(frame_bbs, self.model_image_size,
(im_height, im_width))
frame_dets = []
for box in frame_bbs:
pred_box = self.p1p2Toxywh(box[:4]).data.cpu().numpy()
# box = box.data.cpu().numpy()
cls_conf = box[4].item()
cls_id = box[-1]
cls_name = self.class_names[int(cls_id)]
if classes is not None and cls_name not in classes:
continue
left, top, w, h = pred_box
right = left + w
bottom = top + h
width = right - left + 1
height = bottom - top + 1
width_buffer = width * buffer_ratio
height_buffer = height * buffer_ratio
top = max(0.0, top - 0.5 * height_buffer)
left = max(0.0, left - 0.5 * width_buffer)
bottom = min(im_height - 1.0, bottom + 0.5 * height_buffer)
right = min(im_width - 1.0, right + 0.5 * width_buffer)
box_infos = []
for c in box_format:
if c == 't':
box_infos.append(int(round(top)))
elif c == 'l':
box_infos.append(int(round(left)))
elif c == 'b':
box_infos.append(int(round(bottom)))
elif c == 'r':
box_infos.append(int(round(right)))
elif c == 'w':
box_infos.append(int(round(width + width_buffer)))
elif c == 'h':
box_infos.append(int(round(height + height_buffer)))
else:
assert False, 'box_format given in detect unrecognised!'
assert len(box_infos) > 0, 'box infos is blank'
detection = (box_infos, cls_conf, cls_name)
frame_dets.append(detection)
detections.append(frame_dets)
return detections
@staticmethod
def _resize_boxes(boxes, current_dim, original_shape):
h_ratio = original_shape[0] / current_dim[0]
w_ratio = original_shape[1] / current_dim[1]
boxes[..., 0] *= w_ratio
boxes[..., 1] *= h_ratio
boxes[..., 2] *= w_ratio
boxes[..., 3] *= h_ratio
return boxes
from PIL import Image
import cv2
config_path = 'config/yolov3.cfg'
weights_path = 'config/yolov3.weights'
class_path = 'config/coco.names'
img_size = 416
conf_thres = 0.8
nms_thres = 0.4
# Load model and weights
model = Darknet(config_path, img_size=img_size)
model.load_weights(weights_path)
model.cuda()
model.eval()
classes = load_classes(class_path)
Tensor = torch.cuda.FloatTensor
def detect_image(img):
# Scale and pad image
ratio = min(img_size / img.size[0], img_size / img.size[1])
imw = round(img.size[0] * ratio)
imh = round(img.size[1] * ratio)
img_transforms = transforms.Compose([
transforms.Resize((imh, imw)),
transforms.Pad((max(int(
(imh - imw) / 2), 0), max(int(
(imw - imh) / 2), 0), max(int(
