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 _draw_and_save_output_image(image_path, detections, img_size, output_path, classes):
"""Draws detections in output image and stores this.
:param image_path: Path to input image
:type image_path: str
:param detections: List of detections on image
:type detections: [Tensor]
:param img_size: Size of each image dimension for yolo
:type img_size: int
:param output_path: Path of output directory
:type output_path: str
:param classes: List of class names
:type classes: [str]
"""
# Create plot
img = np.array(Image.open(image_path))
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(img)
# Rescale boxes to original image
detections = rescale_boxes(detections, img_size, img.shape[:2])
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
# Bounding-box colors
cmap = plt.get_cmap("tab20b")
colors = [cmap(i) for i in np.linspace(0, 1, n_cls_preds)]
bbox_colors = random.sample(colors, n_cls_preds)
for x1, y1, x2, y2, conf, cls_pred in detections:
print(f"\t+ Label: {classes[int(cls_pred)]} | Confidence: {conf.item():0.4f}")
box_w = x2 - x1
box_h = y2 - y1
color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor="none")
# Add the bbox to the plot
ax.add_patch(bbox)
# Add label
plt.text(
x1,
y1,
s=classes[int(cls_pred)],
color="white",
verticalalignment="top",
bbox={"color": color, "pad": 0})
# Save generated image with detections
plt.axis("off")
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
filename = os.path.basename(image_path).split(".")[0]
output_path = os.path.join(output_path, f"{filename}.png")
plt.savefig(output_path, bbox_inches="tight", pad_inches=0.0)
plt.close()