yolov3 / utils.py

init model

32865f3 6 months ago

No virus

7.93 kB

	import numpy as np
	import cv2
	import torch
	import time
	import torchvision
	import random


	def box_iou(box1, box2):
	# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
	"""
	Return intersection-over-union (Jaccard index) of boxes.
	Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
	Arguments:
	box1 (Tensor[N, 4])
	box2 (Tensor[M, 4])
	Returns:
	iou (Tensor[N, M]): the NxM matrix containing the pairwise
	IoU values for every element in boxes1 and boxes2
	"""

	def box_area(box):
	# box = 4xn
	return (box[2] - box[0]) * (box[3] - box[1])

	area1 = box_area(box1.T)
	area2 = box_area(box2.T)

	# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
	inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
	return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)


	def plot_one_box(x, image, color=None, label=None, line_thickness=None):
	# Plots one bounding box on image img
	tl = line_thickness or round(
	0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 # line/font thickness
	color = color or [random.randint(0, 255) for _ in range(3)]
	c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
	cv2.rectangle(image, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
	if label:
	tf = max(tl - 1, 1) # font thickness
	t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
	c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
	cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA) # filled
	cv2.putText(image, label, (c1[0], c1[1] - 2), 0, tl / 3,
	[225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)


	def clip_coords(boxes, img_shape):
	# Clip bounding xyxy bounding boxes to image shape (height, width)
	boxes[:, 0].clamp_(0, img_shape[1]) # x1
	boxes[:, 1].clamp_(0, img_shape[0]) # y1
	boxes[:, 2].clamp_(0, img_shape[1]) # x2
	boxes[:, 3].clamp_(0, img_shape[0]) # y2


	def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
	# Rescale coords (xyxy) from img1_shape to img0_shape
	if ratio_pad is None: # calculate from img0_shape
	gain = max(img1_shape) / max(img0_shape) # gain = old / new
	pad = (img1_shape[1] - img0_shape[1] * gain) / \
	2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
	else:
	gain = ratio_pad[0][0]
	pad = ratio_pad[1]

	coords[:, [0, 2]] -= pad[0] # x padding
	coords[:, [1, 3]] -= pad[1] # y padding
	coords[:, :4] /= gain
	clip_coords(coords, img0_shape)
	return coords


	def xywh2xyxy(x):
	# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where
	# xy1=top-left, xy2=bottom-right
	y = torch.zeros_like(x) if isinstance(
	x, torch.Tensor) else np.zeros_like(x)
	y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
	y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
	y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
	y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
	return y


	def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=True,
	scaleFill=False, scaleup=True):
	# Resize image to a 32-pixel-multiple rectangle
	# https://github.com/ultralytics/yolov3/issues/232
	shape = img.shape[:2] # current shape [height, width]
	if isinstance(new_shape, int):
	new_shape = (new_shape, new_shape)

	# Scale ratio (new / old)
	r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
	if not scaleup: # only scale down, do not scale up (for better test mAP)
	r = min(r, 1.0)

	# Compute padding
	ratio = r, r # width, height ratios
	new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
	dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - \
	new_unpad[1] # wh padding
	if auto: # minimum rectangle
	dw, dh = np.mod(dw, 32), np.mod(dh, 32) # wh padding
	elif scaleFill: # stretch
	dw, dh = 0.0, 0.0
	new_unpad = new_shape
	ratio = new_shape[0] / shape[1], new_shape[1] / \
	shape[0] # width, height ratios

	dw /= 2 # divide padding into 2 sides
	dh /= 2

	if shape[::-1] != new_unpad: # resize
	img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
	top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
	left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
	img = cv2.copyMakeBorder(img, top, bottom, left, right,
	cv2.BORDER_CONSTANT, value=color) # add border
	return img, ratio, (dw, dh)


	def non_max_suppression(
	prediction,
	conf_thres=0.1,
	iou_thres=0.6,
	multi_label=True,
	classes=None,
	agnostic=False):
	"""
	Performs Non-Maximum Suppression on inference results
	Returns detections with shape:
	nx6 (x1, y1, x2, y2, conf, cls)
	"""

	# Settings
	merge = True # merge for best mAP
	# (pixels) minimum and maximum box width and height
	min_wh, max_wh = 2, 4096
	time_limit = 10.0 # seconds to quit after

	t = time.time()
	nc = prediction[0].shape[1] - 5 # number of classes
	multi_label &= nc > 1 # multiple labels per box
	output = [None] * prediction.shape[0]
	for xi, x in enumerate(prediction): # image index, image inference
	# Apply constraints
	x = x[x[:, 4] > conf_thres] # confidence
	x = x[((x[:, 2:4] > min_wh) & (x[:, 2:4] < max_wh)).all(1)]

	# If none remain process next image
	if not x.shape[0]:
	continue

	# Compute conf
	x[..., 5:] = x[..., 4:5] # conf = obj_conf cls_conf

	# Box (center x, center y, width, height) to (x1, y1, x2, y2)
	box = xywh2xyxy(x[:, :4])

	# Detections matrix nx6 (xyxy, conf, cls)
	if multi_label:
	i, j = (x[:, 5:] > conf_thres).nonzero().t()
	x = torch.cat((box[i], x[i, j + 5].unsqueeze(1),
	j.float().unsqueeze(1)), 1)
	else: # best class only
	conf, j = x[:, 5:].max(1)
	x = torch.cat(
	(box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[
	conf > conf_thres]

	# Filter by class
	if classes:
	x = x[(j.view(-1, 1) == torch.tensor(classes,
	device=j.device)).any(1)]

	# Apply finite constraint
	# if not torch.isfinite(x).all():
	# x = x[torch.isfinite(x).all(1)]

	# If none remain process next image
	n = x.shape[0] # number of boxes
	if not n:
	continue

	# Sort by confidence
	# x = x[x[:, 4].argsort(descending=True)]

	# Batched NMS
	c = x[:, 5] * 0 if agnostic else x[:, 5] # classes
	boxes, scores = x[:, :4].clone() + c.view(-1, 1) * \
	max_wh, x[:, 4] # boxes (offset by class), scores
	i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
	if merge and (
	1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
	try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
	iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
	weights = iou * scores[None] # box weights
	x[i, :4] = torch.mm(weights, x[:, :4]).float(
	) / weights.sum(1, keepdim=True) # merged boxes
	# i = i[iou.sum(1) > 1] # require redundancy
	except BaseException:
	# https://github.com/ultralytics/yolov3/issues/1139
	# print(x, i, x.shape, i.shape)
	pass

	output[xi] = x[i]
	if (time.time() - t) > time_limit:
	break # time limit exceeded

	return output