def create_dataset(opt, model_cfg):
data_config = parse_data_config(opt.data)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
train_loader, valid_loader, test_loader = None, None, None
train_dataset = ListDataset(train_path,
img_size=int(model_cfg[0]['width']),
augment=True,
multiscale=opt.multiscale_training)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True,
collate_fn=train_dataset.collate_fn,
)
#import pdb;pdb.set_trace()
valid_dataset = ListDataset(train_path,
img_size=int(model_cfg[0]['width']),
augment=False,
multiscale=False)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=4,
collate_fn=valid_dataset.collate_fn)
test_loader = valid_loader
return class_names, train_loader, valid_loader, test_loader
def __init__(self, opt):
data_config = parse_data_config(opt.data)
backup = data_config["backup"]
if not os.path.exists(backup): os.makedirs(backup)
model_name = os.path.splitext(os.path.split(
opt.cfg)[-1])[0] + '_best.pth'
self.save_path = os.path.join(backup, model_name + '_best.pth')
self.best_map = 0
def get_data(opt, num_imgs, n_cpu):
'''Get data'''
data_cfg = parse_data_config(opt.data_config_path)
dataset_path = data_cfg['train']
if opt.dataset == 'coco':
opt.img_size = 416
Dataname = ListDataset
elif opt.dataset == 'imagenet':
opt.img_size = 224
Dataname = ImagenetDataset
dataset = Dataname(dataset_path, img_size=opt.img_size)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=num_imgs,
shuffle=True,
num_workers=n_cpu,
pin_memory=True)
return dataset, dataloader
def get_para(opt):
'''Training paras'''
parameters = Paras(opt.dataset)
# Get data path
data_config = parse_data_config(opt.data_config_path)
if opt.dataset == 'dota':
parameters.data_path = '/home/cheney/data/dota/train_crop.lmdb'
elif opt.dataset == 'nwpuvhr':
parameters.data_path = data_config['train']
# Get hyper parameters
# hyperparams = parse_model_config(opt.model_config_path)[0]
# para.epochs = int(hyperparams['max_epoch'])
# para.lr = float(hyperparams['learning_rate'])
# para.lr_steps = list(map(int, hyperparams['steps'].strip().split(',')))
# lr_scales = map(float, hyperparams['scales'].strip().split(','))
# momentum = float(hyperparams['momentum'])
# decay = float(hyperparams['decay'])
return parameters
def __init__(self, data_config, net_config, weights, image_size):
"""Class init function."""
QtCore.QThread.__init__(self)
self.image_list = []
self.threshold = 0.9
self.image_directory = ''
self.data = None
self.image_size = image_size
if torch.cuda.is_available():
self.device = torch.device('cuda:0')
else:
self.device = torch.device('cpu')
self.data_config = parse_data_config(data_config)
# Extracts class labels from file
self.classes = load_classes(self.data_config['names'])
self.model = Darknet(net_config, image_size)
checkpoint = torch.load(weights, map_location='cpu')
self.model.load_state_dict(checkpoint['model'])
def train(opt, model, train_loader, optimizer, epoch=0):
model.train()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
for batch_i, (_, imgs, targets) in enumerate(train_loader):
start_time = time.time()
batches_done = len(train_loader) * epoch + batch_i
imgs = Variable(imgs.to(device))
targets = Variable(targets.to(device), requires_grad=False)
#import pdb;pdb.set_trace()
optimizer.zero_grad()
loss, outputs = model(imgs, targets)
loss.backward()
optimizer.step()
# ----------------
# Log progress
# ----------------
time_left = datetime.timedelta(seconds=time.time() - start_time)
log_str = "[Epoch %d/%d, Batch %d/%d] loss:%f " % (
epoch, opt.epochs, batch_i, len(train_loader) // opt.batch_size,
loss)
log_str += f"{time_left} "
log_str += "cls_acc:%f " % model.yolo_layers[-1].metrics["cls_acc"]
log_str += "recall50:%f " % model.yolo_layers[-1].metrics["recall50"]
if batch_i % 5 == 0:
print(log_str)
model.seen += imgs.size(0)
#save
if epoch % 1 == 0:
data_config = parse_data_config(opt.data)
backup = data_config["backup"]
if not os.path.exists(backup): os.makedirs(backup)
model_name = os.path.splitext(os.path.split(opt.cfg)[-1])[0]
save_path = os.path.join(backup, model_name + '_last.pth')
if os.path.exists(save_path): os.remove(save_path)
torch.save(model.state_dict(), save_path)
def run():
print_environment_info()
parser = argparse.ArgumentParser(description="Evaluate validation data.")
parser.add_argument("-m",
"--model",
type=str,
default="config/yolov3.cfg",
help="Path to model definition file (.cfg)")
parser.add_argument(
"-w",
"--weights",
type=str,
default="weights/yolov3.weights",
help="Path to weights or checkpoint file (.weights or .pth)")
parser.add_argument("-d",
"--data",
type=str,
default="config/coco.data",
help="Path to data config file (.data)")
parser.add_argument("-b",
"--batch_size",
type=int,
default=8,
help="Size of each image batch")
parser.add_argument("-v",
"--verbose",
action='store_true',
help="Makes the validation more verbose")
parser.add_argument("--img_size",
type=int,
default=416,
help="Size of each image dimension for yolo")
parser.add_argument(
"--n_cpu",
type=int,
default=8,
help="Number of cpu threads to use during batch generation")
parser.add_argument("--iou_thres",
type=float,
default=0.5,
help="IOU threshold required to qualify as detected")
parser.add_argument("--conf_thres",
type=float,
default=0.01,
help="Object confidence threshold")
parser.add_argument("--nms_thres",
type=float,
default=0.4,
help="IOU threshold for non-maximum suppression")
args = parser.parse_args()
print("Command line arguments: {}".format(args))
# Load configuration from data file
data_config = parse_data_config(args.data)
# Path to file containing all images for validation
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"]) # List of class names
precision, recall, AP, f1, ap_class = evaluate_model_file(
args.model,
args.weights,
valid_path,
class_names,
batch_size=args.batch_size,
img_size=args.img_size,
n_cpu=args.n_cpu,
iou_thres=args.iou_thres,
conf_thres=args.conf_thres,
nms_thres=args.nms_thres,
verbose=True)
if re.findall('([-\w]+\.(?:jpg|gif|png))',
os.path.basename(file_path[0].lower())):
return render_template('person_reid.html')
else:
return render_template('person_det.html')
print('Load Input Arguments')
args = parse_args()
print('Load Object Detection model ...')
person_handler = PersonHandler(args)
print('Load Label Map')
cls_dict = load_cls_dict(parse_data_config(args.data_path)['names'])
# grab image and do object detection (until stopped by user)
print('starting to loop and detect')
vis = BBoxVisualization(cls_dict)
@app.route('/video_feed', methods=['GET'])
def video_feed():
global loader
loader = LoadCamera(file_path[0], args.img_size)
return Response(person_handler.loop_and_detect(loader, vis),
mimetype='multipart/x-mixed-replace; boundary=frame')
if __name__ == '__main__':
def run():
print_environment_info()
parser = argparse.ArgumentParser(description="Trains the YOLO model.")
parser.add_argument("-m", "--model", type=str, default="config/yolov3.cfg", help="Path to model definition file (.cfg)")
parser.add_argument("-d", "--data", type=str, default="config/coco.data", help="Path to data config file (.data)")
parser.add_argument("-e", "--epochs", type=int, default=300, help="Number of epochs")
parser.add_argument("-v", "--verbose", action='store_true', help="Makes the training more verbose")
parser.add_argument("--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation")
parser.add_argument("--pretrained_weights", type=str, help="Path to checkpoint file (.weights or .pth). Starts training from checkpoint model")
parser.add_argument("--checkpoint_interval", type=int, default=1, help="Interval of epochs between saving model weights")
parser.add_argument("--evaluation_interval", type=int, default=1, help="Interval of epochs between evaluations on validation set")
parser.add_argument("--multiscale_training", action="store_false", help="Allow for multi-scale training")
parser.add_argument("--iou_thres", type=float, default=0.5, help="Evaluation: IOU threshold required to qualify as detected")
parser.add_argument("--conf_thres", type=float, default=0.1, help="Evaluation: Object confidence threshold")
parser.add_argument("--nms_thres", type=float, default=0.5, help="Evaluation: IOU threshold for non-maximum suppression")
parser.add_argument("--logdir", type=str, default="logs", help="Directory for training log files (e.g. for TensorBoard)")
args = parser.parse_args()
print("Command line arguments: {}".format(args))
logger = Logger(args.logdir) # Tensorboard logger
# Create output directories if missing
os.makedirs("output", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)
# Get data configuration
data_config = parse_data_config(args.data)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ############
# Create model
# ############
model = load_model(args.model, args.pretrained_weights)
# Print model
if args.verbose:
summary(model, input_size=(3, model.hyperparams['height'], model.hyperparams['height']))
mini_batch_size = model.hyperparams['batch'] // model.hyperparams['subdivisions']
# #################
# Create Dataloader
# #################
# Load training dataloader
dataloader = _create_data_loader(
train_path,
mini_batch_size,
model.hyperparams['height'],
args.n_cpu,
args.multiscale_training)
# Load validation dataloader
validation_dataloader = _create_validation_data_loader(
valid_path,
mini_batch_size,
model.hyperparams['height'],
args.n_cpu)
# ################
# Create optimizer
# ################
params = [p for p in model.parameters() if p.requires_grad]
if (model.hyperparams['optimizer'] in [None, "adam"]):
optimizer = optim.Adam(
params,
lr=model.hyperparams['learning_rate'],
weight_decay=model.hyperparams['decay'],
)
elif (model.hyperparams['optimizer'] == "sgd"):
optimizer = optim.SGD(
params,
lr=model.hyperparams['learning_rate'],
weight_decay=model.hyperparams['decay'],
momentum=model.hyperparams['momentum'])
else:
print("Unknown optimizer. Please choose between (adam, sgd).")
for epoch in range(args.epochs):
print("\n---- Training Model ----")
model.train() # Set model to training mode
for batch_i, (_, imgs, targets) in enumerate(tqdm.tqdm(dataloader, desc="Training Epoch {}".format(epoch))):
batches_done = len(dataloader) * epoch + batch_i
imgs = imgs.to(device, non_blocking=True)
targets = targets.to(device)
outputs = model(imgs)
loss, loss_components = compute_loss(outputs, targets, model)
loss.backward()
###############
# Run optimizer
###############
if batches_done % model.hyperparams['subdivisions'] == 0:
# Adapt learning rate
# Get learning rate defined in cfg
lr = model.hyperparams['learning_rate']
if batches_done < model.hyperparams['burn_in']:
# Burn in
lr *= (batches_done / model.hyperparams['burn_in'])
else:
# Set and parse the learning rate to the steps defined in the cfg
for threshold, value in model.hyperparams['lr_steps']:
if batches_done > threshold:
lr *= value
# Log the learning rate
logger.scalar_summary("train/learning_rate", lr, batches_done)
# Set learning rate
for g in optimizer.param_groups:
g['lr'] = lr
# Run optimizer
optimizer.step()
# Reset gradients
optimizer.zero_grad()
# ############
# Log progress
# ############
if args.verbose:
print(AsciiTable(
[
["Type", "Value"],
["IoU loss", float(loss_components[0])],
["Object loss", float(loss_components[1])],
["Class loss", float(loss_components[2])],
["Loss", float(loss_components[3])],
["Batch loss", to_cpu(loss).item()],
]).table)
# Tensorboard logging
tensorboard_log = [
("train/iou_loss", float(loss_components[0])),
("train/obj_loss", float(loss_components[1])),
("train/class_loss", float(loss_components[2])),
("train/loss", to_cpu(loss).item())]
logger.list_of_scalars_summary(tensorboard_log, batches_done)
model.seen += imgs.size(0)
# #############
# Save progress
# #############
# Save model to checkpoint file
if epoch % args.checkpoint_interval == 0:
checkpoint_path = "checkpoints/yolov3_ckpt_{}.pth".format(epoch)
print("---- Saving checkpoint to: '{}' ----".format(checkpoint_path))
torch.save(model.state_dict(), checkpoint_path)
# ########
# Evaluate
# ########
if epoch % args.evaluation_interval == 0:
print("\n---- Evaluating Model ----")
# Evaluate the model on the validation set
metrics_output = _evaluate(
model,
validation_dataloader,
class_names,
img_size=model.hyperparams['height'],
iou_thres=args.iou_thres,
conf_thres=args.conf_thres,
nms_thres=args.nms_thres,
verbose=args.verbose
)
if metrics_output is not None:
precision, recall, AP, f1, ap_class = metrics_output
evaluation_metrics = [
("validation/precision", precision.mean()),
("validation/recall", recall.mean()),
("validation/mAP", AP.mean()),
("validation/f1", f1.mean())]
logger.list_of_scalars_summary(evaluation_metrics, epoch)
# print(new_weights)
# save new weights
weighs_file = './4_sparsity_merged.weights'
if SAVE:
torch.save(new_weights, weighs_file)
import argparse
import time
from model import *
from utils.datasets import *
from utils.utils import *
data_config_path = 'cfg/coco.data'
data_config = parse_config.parse_data_config(data_config_path)
images_path = './data/samples'
batch_size = 1
conf_thres = 0.25
nms_thres = 0.45
new_model = Darknet_new(net_config_path, img_size)
new_model.load_weights(weighs_file)
new_model.to(device).eval()
# Set Dataloader
classes = load_classes(data_config['names']) # Extracts class labels from file
dataloader = load_images(images_path, batch_size=batch_size, img_size=img_size)
imgs = [] # Stores image paths
img_detections = [] # Stores detections for each image index
def train(
cfg,
data_cfg,
img_size=416,
resume=False,
epochs=270,
batch_size=16,
accumulate=1,
multi_scale=False,
freeze_backbone=False,
num_workers=4,
transfer=False, # Transfer learning (train only YOLO layers)
use_cpu=True,
backend='nccl',
world_size=1,
rank=0,
dist_url='tcp://127.0.0.1:9999'):
weights = 'weights' + os.sep
latest = weights + 'latest.pt'
best = weights + 'best.pt'
use_gpu = torch.cuda.is_available()
if use_cpu:
use_gpu = False
device = torch.device('cuda' if use_gpu else 'cpu')
if multi_scale:
img_size = 608 # initiate with maximum multi_scale size
num_workers = 0 # bug https://github.com/ultralytics/yolov3/issues/174
else:
torch.backends.cudnn.benchmark = True # unsuitable for multiscale
# Configure run
train_path = parse_data_config(data_cfg)['train']
# Initialize model
model = Darknet(cfg, img_size, device).to(device)
# Optimizer
lr0 = 0.001 # initial learning rate
optimizer = torch.optim.SGD(model.parameters(),
lr=lr0,
momentum=0.9,
weight_decay=0.0005)
cutoff = -1 # backbone reaches to cutoff layer
start_epoch = 0
best_loss = float('inf')
yl = get_yolo_layers(model) # yolo layers
nf = int(model.module_defs[yl[0] -
1]['filters']) # yolo layer size (i.e. 255)
if resume: # Load previously saved model
if transfer: # Transfer learning
chkpt = torch.load(weights + 'yolov3-spp.pt', map_location=device)
model.load_state_dict(
{
k: v
for k, v in chkpt['model'].items()
if v.numel() > 1 and v.shape[0] != 255
},
strict=False)
for p in model.parameters():
p.requires_grad = True if p.shape[0] == nf else False
else: # resume from latest.pt
chkpt = torch.load(latest, map_location=device) # load checkpoint
model.load_state_dict(chkpt['model'])
start_epoch = chkpt['epoch'] + 1
if chkpt['optimizer'] is not None:
optimizer.load_state_dict(chkpt['optimizer'])
best_loss = chkpt['best_loss']
del chkpt
else: # Initialize model with backbone (optional)
if '-tiny.cfg' in cfg:
cutoff = model.load_darknet_weights(weights +
'yolov3-tiny.conv.15')
else:
cutoff = model.load_darknet_weights(weights + 'darknet53.conv.74')
# Set scheduler (reduce lr at epoch 250)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[250],
gamma=0.1,
last_epoch=start_epoch -
1)
# Dataset
dataset = LoadImagesAndLabels(train_path, img_size=img_size, augment=True)
# Initialize distributed training
if torch.cuda.device_count() > 1:
dist.init_process_group(backend=backend,
init_method=dist_url,
world_size=world_size,
rank=rank)
model = torch.nn.parallel.DistributedDataParallel(model)
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
else:
sampler = None
# Dataloader
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False,
pin_memory=False,
collate_fn=dataset.collate_fn,
sampler=sampler)
# Start training
t = time.time()
model_info(model)
nB = len(dataloader)
n_burnin = min(round(nB / 5 + 1), 1000) # burn-in batches
for epoch in range(start_epoch, epochs):
model.train()
print(
('\n%8s%12s' + '%10s' * 7) % ('Epoch', 'Batch', 'xy', 'wh', 'conf',
'cls', 'total', 'nTargets', 'time'))
# Update scheduler
scheduler.step()
# Freeze backbone at epoch 0, unfreeze at epoch 1
if freeze_backbone and epoch < 2:
for name, p in model.named_parameters():
if int(name.split('.')[1]) < cutoff: # if layer < 75
p.requires_grad = False if epoch == 0 else True
mloss = defaultdict(float) # mean loss
for i, (imgs, targets, _, _) in enumerate(dataloader):
imgs = imgs.to(device)
targets = targets.to(device)
nT = len(targets)
if nT == 0: # if no targets continue
continue
# Plot images with bounding boxes
plot_images = False
if plot_images:
fig = plt.figure(figsize=(10, 10))
for ip in range(len(imgs)):
boxes = xywh2xyxy(targets[targets[:, 0] == ip,
2:6]).numpy().T * img_size
plt.subplot(4, 4,
ip + 1).imshow(imgs[ip].numpy().transpose(
1, 2, 0))
plt.plot(boxes[[0, 2, 2, 0, 0]], boxes[[1, 1, 3, 3, 1]],
'.-')
plt.axis('off')
fig.tight_layout()
fig.savefig('batch_%g.jpg' % i, dpi=fig.dpi)
# SGD burn-in
if epoch == 0 and i <= n_burnin:
lr = lr0 * (i / n_burnin)**4
for x in optimizer.param_groups:
x['lr'] = lr
# Run model
pred = model(imgs)
# Build targets
target_list = build_targets(model, targets)
# Compute loss
loss, loss_dict = compute_loss(pred, target_list)
# Compute gradient
loss.backward()
# Accumulate gradient for x batches before optimizing
if (i + 1) % accumulate == 0 or (i + 1) == nB:
optimizer.step()
optimizer.zero_grad()
# Running epoch-means of tracked metrics
for key, val in loss_dict.items():
mloss[key] = (mloss[key] * i + val) / (i + 1)
s = ('%8s%12s' + '%10.3g' * 7) % (
'%g/%g' % (epoch, epochs - 1), '%g/%g' %
(i, nB - 1), mloss['xy'], mloss['wh'], mloss['conf'],
mloss['cls'], mloss['total'], nT, time.time() - t)
t = time.time()
print(s)
# Multi-Scale training (320 - 608 pixels) every 10 batches
if multi_scale and (i + 1) % 10 == 0:
dataset.img_size = random.choice(range(10, 20)) * 32
print('multi_scale img_size = %g' % dataset.img_size)
# Update best loss
if mloss['total'] < best_loss:
best_loss = mloss['total']
# Save training results
save = True
if save:
# Save latest checkpoint
chkpt = {
'epoch':
epoch,
'best_loss':
best_loss,
'model':
model.module.state_dict()
if type(model) is nn.parallel.DistributedDataParallel else
model.state_dict(),
'optimizer':
optimizer.state_dict()
}
torch.save(chkpt, latest)
# Save best checkpoint
if best_loss == mloss['total']:
torch.save(chkpt, best)
# Save backup every 10 epochs (optional)
if epoch > 0 and epoch % 10 == 0:
torch.save(chkpt, weights + 'backup%g.pt' % epoch)
del chkpt
# Calculate mAP
with torch.no_grad():
results = test.test(cfg,
data_cfg,
batch_size=batch_size,
img_size=img_size,
model=model)
# Write epoch results
with open('results.txt', 'a') as file:
file.write(s + '%11.3g' * 3 % results + '\n') # append P, R, mAP
def test(cfg,
data_cfg,
weights=None,
batch_size=16,
img_size=416,
iou_thres=0.5,
conf_thres=0.1,
nms_thres=0.5,
save_json=False,
use_cpu=True,
model=None):
if model is None:
use_gpu = torch.cuda.is_available()
if use_cpu:
use_gpu = False
device = torch.device('cuda' if use_gpu else 'cpu')
# Initialize model
model = Darknet(cfg, img_size, device)
# Load weights
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(
torch.load(weights, map_location=device)['model'])
else: # darknet format
model.load_darknet_weights(weights)
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
else:
device = next(model.parameters()).device # get model device
# Configure run
data_cfg = parse_data_config(data_cfg)
test_path = data_cfg['valid']
# Dataloader
dataset = LoadImagesAndLabels(test_path, img_size=img_size)
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=4,
pin_memory=False,
collate_fn=dataset.collate_fn)
model.eval()
seen = 0
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP'))
mP, mR, mAP, mAPj = 0.0, 0.0, 0.0, 0.0
jdict, tdict, stats, AP, AP_class = [], [], [], [], []
coco91class = coco80_to_coco91_class()
for batch_i, (imgs, targets, paths,
shapes) in enumerate(tqdm(dataloader,
desc='Calculating mAP')):
targets = targets.to(device)
imgs = imgs.to(device)
output = model(imgs)
output = non_max_suppression(output,
conf_thres=conf_thres,
nms_thres=nms_thres)
# Per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
correct, detected = [], []
tcls = torch.Tensor()
seen += 1
if pred is None:
if len(labels):
tcls = labels[:, 0].cpu() # target classes
stats.append(
(correct, torch.Tensor(), torch.Tensor(), tcls))
continue
if save_json: # add to json pred dictionary
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(Path(paths[si]).stem.split('_')[-1])
box = pred[:, :4].clone() # xyxy
scale_coords(img_size, box, shapes[si]) # to original shape
box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for di, d in enumerate(pred):
jdict.append({
'image_id': image_id,
'category_id': coco91class[int(d[6])],
'bbox': [float3(x) for x in box[di]],
'score': float(d[4])
})
if len(labels):
# Extract target boxes as (x1, y1, x2, y2)
tbox = xywh2xyxy(labels[:, 1:5]) * img_size # target boxes
tcls = labels[:, 0] # target classes
for *pbox, pconf, pcls_conf, pcls in pred:
if pcls not in tcls:
correct.append(0)
continue
# Best iou, index between pred and targets
iou, bi = bbox_iou(pbox, tbox).max(0)
# If iou > threshold and class is correct mark as correct
if iou > iou_thres and bi not in detected:
correct.append(1)
detected.append(bi)
else:
correct.append(0)
else:
# If no labels add number of detections as incorrect
correct.extend([0] * len(pred))
# Append Statistics (correct, conf, pcls, tcls)
stats.append(
(correct, pred[:, 4].cpu(), pred[:, 6].cpu(), tcls.cpu()))
# Compute means
stats_np = [np.concatenate(x, 0) for x in list(zip(*stats))]
if len(stats_np):
AP, AP_class, R, P = ap_per_class(*stats_np)
mP, mR, mAP = P.mean(), R.mean(), AP.mean()
# Print P, R, mAP
print(('%11s%11s' + '%11.3g' * 3) % (seen, len(dataset), mP, mR, mAP))
# Print mAP per class
if len(stats_np):
print('\nmAP Per Class:')
names = load_classes(data_cfg['names'])
for c, a in zip(AP_class, AP):
print('%15s: %-.4f' % (names[c], a))
# Save JSON
if save_json and mAP and len(jdict):
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataset.img_files]
with open('results.json', 'w') as file:
json.dump(jdict, file)
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
# https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
cocoGt = COCO('../coco/annotations/instances_val2014.json'
) # initialize COCO ground truth api
cocoDt = cocoGt.loadRes('results.json') # initialize COCO pred api
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
cocoEval.params.imgIds = imgIds # [:32] # only evaluate these images
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
mAP = cocoEval.stats[1] # update mAP to pycocotools mAP
# Return mAP
return mP, mR, mAP
type=int,
default=100,
help="metric printing interval")
opt = parser.parse_args()
print(opt)
for dir_name in ["output", "logs", "trained_models", "tensorboard_logs"]:
os.makedirs(dir_name, exist_ok=True)
tensorboard_logger = Logger("tensorboard_logs")
date_str = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
log_filename = date_str + "_train_logs.log"
log_save_path = "logs/" + log_filename
logger = create_logger(log_save_path=log_save_path)
"""Get data configuration"""
data_config = parse_data_config(path=opt.data_config)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
trained_models_dir = "trained_models/" + date_str + "/"
"""Initiate model"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
"""If specified we start from checkpoint"""
darknet_pretrained = False
if opt.pretrained_weights:
print("load pretrained weights")
if opt.pretrained_weights.endswith(".pth"):
"""Load darknet weights"""
model.load_state_dict(torch.load(opt.pretrained_weights))
from __future__ import division
from utils.utils import load_classes, printBBoxes
from utils.datasets import TestDataset
from utils.parse_config import parse_data_config
import torch
if __name__ == "__main__":
epochs = 100
model_def = 'config/yolov3visdrone.cfg'
data_config = 'config/visdrone.data'
data_config = parse_data_config(data_config)
class_names = load_classes(data_config["names"])
train_path = data_config["train"]
valid_path = data_config["valid"]
# Get dataloader
dataset = TestDataset(train_path)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=1,
num_workers=0, # opt.n_cpu,
pin_memory=True)
for epoch in range(epochs):
for batch_i, (paths, targets) in enumerate(dataloader):
printBBoxes(paths[0], targets[0], class_names, rescale=False)
if __name__ == "__main__":
torch.manual_seed(0)
np.random.seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("1. Path to save the output")
args.save_path = Path(args.save_path)
args.save_path.makedirs_p()
print("=> will save everything to {}".format(args.save_path))
print("2. Data Loading...")
if args.is_gt:
data_config = parse_data_config(args.data_config)
test_path = data_config["test"]
class_names = load_classes(
data_config["names"]
) # list of all class names，in the same numbering order as the annotaion
test_set = ListDataset(test_path,
transform=None,
img_size=args.img_size,
multiscale=False)
test_loader = DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_cpu,
collate_fn=test_set.collate_fn)
else:
test_set = ImageFolder(args.image_folder, img_size=args.img_size)
# Hacky overide of defaults
opt.model_def = 'config/yolov3visdrone.cfg'
opt.data_config = 'config/visdrone.data'
opt.multiscale_training = False
print(opt)
logger = Logger("logs/exp" + str(len(os.listdir('./logs')) + 1))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("output", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)
# Get data configuration
data_config = parse_data_config(opt.data_config)
if platform == "linux" or platform == "linux2":
train_path = data_config["train_Linux"]
valid_path = data_config["valid_Linux"]
else:
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
# Initiate model
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
def detect(
cfg,
data_cfg,
weights,
images,
output='output', # output folder
img_size=416,
conf_thres=0.5,
nms_thres=0.5,
save_images=True,
use_cpu=True):
use_gpu = torch.cuda.is_available()
if use_cpu:
use_gpu = False
device = torch.device('cuda' if use_gpu else 'cpu')
if os.path.exists(output):
shutil.rmtree(output) # delete output folder
os.makedirs(output) # make new output folder
# Initialize model
model = Darknet(cfg, img_size, device)
# Load weights
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(
torch.load(weights, map_location=device)['model'])
else: # darknet format
model.load_darknet_weights(weights)
model.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
dataloader = LoadImages(images, img_size=img_size)
# Get classes and colors
print('Load Label Map')
cls_dict = load_cls_dict(parse_data_config(data_cfg)['names'])
# grab image and do object detection (until stopped by user)
print('starting to loop and detect')
vis = BBoxVisualization(cls_dict)
for i, (path, img, im0, vid_cap) in enumerate(dataloader):
t = time.time()
save_path = str(Path(output) / Path(path).name)
# Get detections
img = torch.from_numpy(img).float().unsqueeze(0).to(device)
pred = model(img)
detections = non_max_suppression(pred, conf_thres, nms_thres)[0]
if detections is not None and len(detections) > 0:
box, conf, cls = boxes_filtering(im0, detections, img_size)
vis_box = []
for i in range(len(box)):
x1, y1, w, h = box[i]
vis_box.append([y1, x1, y1 + h, x1 + w])
im0 = vis.draw_bboxes(im0, vis_box, conf, cls)
print('Done. (%.3fs)' % (time.time() - t))
if save_images: # Save generated image with detections
if dataloader.mode == 'video':
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
width = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = vid_cap.get(cv2.CAP_PROP_FPS)
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*'avc1'), fps,
(width, height))
vid_writer.write(im0)
else:
cv2.imwrite(save_path, im0)
if save_images and platform == 'darwin': # macos
os.system('open ' + output + ' ' + save_path)
