def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5):
"""Inferences one image with model.
:param model: Model for inference
:type model: models.Darknet
:param image: Image to inference
:type image: nd.array
:param img_size: Size of each image dimension for yolo, defaults to 416
:type img_size: int, optional
:param conf_thres: Object confidence threshold, defaults to 0.5
:type conf_thres: float, optional
:param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
:type nms_thres: float, optional
:return: Detections on image with each detection in the format: [x1, y1, x2, y2, confidence, class]
:rtype: nd.array
"""
model.eval() # Set model to evaluation mode
# Configure input
input_img = transforms.Compose([DEFAULT_TRANSFORMS,
Resize(img_size)])((image, np.zeros(
(1, 5))))[0].unsqueeze(0)
if torch.cuda.is_available():
input_img = input_img.to("cuda")
# Get detections
with torch.no_grad():
detections = model(input_img)
detections = non_max_suppression(detections, conf_thres, nms_thres)
detections = rescale_boxes(detections[0], img_size, image.shape[:2])
return to_cpu(detections).numpy()
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)")
parser.add_argument("--seed",
type=int,
default=-1,
help="Makes results reproducable. Set -1 to disable.")
args = parser.parse_args()
print(f"Command line arguments: {args}")
if args.seed != -1:
provide_determinism(args.seed)
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=f"Training Epoch {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 = f"checkpoints/yolov3_ckpt_{epoch}.pth"
print(f"---- Saving checkpoint to: '{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)
def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres,
nms_thres, verbose):
"""Evaluate model on validation dataset.
:param model: Model to evaluate
:type model: models.Darknet
:param dataloader: Dataloader provides the batches of images with targets
:type dataloader: DataLoader
:param class_names: List of class names
:type class_names: [str]
:param img_size: Size of each image dimension for yolo
:type img_size: int
:param iou_thres: IOU threshold required to qualify as detected
:type iou_thres: float
:param conf_thres: Object confidence threshold
:type conf_thres: float
:param nms_thres: IOU threshold for non-maximum suppression
:type nms_thres: float
:param verbose: If True, prints stats of model
:type verbose: bool
:return: Returns precision, recall, AP, f1, ap_class
"""
model.eval() # Set model to evaluation mode
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for _, imgs, targets in tqdm.tqdm(dataloader, desc="Validating"):
# Extract labels
labels += targets[:, 1].tolist()
# Rescale target
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= img_size
imgs = Variable(imgs.type(Tensor), requires_grad=False)
with torch.no_grad():
outputs = to_cpu(model(imgs))
outputs = non_max_suppression(outputs,
conf_thres=conf_thres,
iou_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs,
targets,
iou_threshold=iou_thres)
if len(sample_metrics) == 0: # No detections over whole validation set.
print("---- No detections over whole validation set ----")
return None
# Concatenate sample statistics
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_metrics))
]
metrics_output = ap_per_class(true_positives, pred_scores, pred_labels,
labels)
print_eval_stats(metrics_output, class_names, verbose)
return metrics_output