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 detect(model, dataloader, output_path, img_size, conf_thres, nms_thres):
"""Inferences images with model.
:param model: Model for inference
:type model: models.Darknet
:param dataloader: Dataloader provides the batches of images to inference
:type dataloader: DataLoader
:param output_path: Path to output directory
:type output_path: str
: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: List of detections. The coordinates are given for the padded image that is provided by the dataloader.
Use `utils.rescale_boxes` to transform them into the desired input image coordinate system before its transformed by the dataloader),
List of input image paths
:rtype: [Tensor], [str]
"""
# Create output directory, if missing
os.makedirs(output_path, exist_ok=True)
model.eval() # Set model to evaluation mode
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
img_detections = [] # Stores detections for each image index
imgs = [] # Stores image paths
for (img_paths, input_imgs) in tqdm.tqdm(dataloader, desc="Detecting"):
# Configure input
input_imgs = Variable(input_imgs.type(Tensor))
# Get detections
with torch.no_grad():
detections = model(input_imgs)
detections = non_max_suppression(detections, conf_thres, nms_thres)
# Store image and detections
img_detections.extend(detections)
imgs.extend(img_paths)
return img_detections, imgs
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 = 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