학습 부분이 끝나면 다음 부분은 감지 부분입니다.
컨벤션 룩 구조
VOC_CLASS_BGR은 쉽게 구별하기 위해 다른 범주에 사용해야 하는 색상 프레임입니다. 예를 들어 빨간색은 A에 사용되고 녹색은 B에 사용되며 도중에 색상을 혼합하기가 쉽지 않습니다.
사진 프레임
def visualize_boxes(image_bgr, boxes, class_names, probs, name_bgr_dict=None, line_thickness=2):
if name_bgr_dict is None:
name_bgr_dict = VOC_CLASS_BGR
image_boxes = image_bgr.copy()#分配到新内存中去
for box, class_name, prob in zip(boxes, class_names, probs):
# Draw box on the image.
left_top, right_bottom = box
left, top = int(left_top[0]), int(left_top[1])
right, bottom = int(right_bottom[0]), int(right_bottom[1])
bgr = name_bgr_dict[class_name]
cv2.rectangle(image_boxes, (left, top), (right, bottom), bgr, thickness=line_thickness)
# Draw text on the image.
text = '%s %.2f' % (class_name, prob)
size, baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=2)
text_w, text_h = size
x, y = left, top
x1y1 = (x, y)
x2y2 = (x + text_w + line_thickness, y + text_h + line_thickness + baseline)
cv2.rectangle(image_boxes, x1y1, x2y2, bgr, -1)
cv2.putText(image_boxes, text, (x + line_thickness, y + 2*baseline + line_thickness),
cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, color=(255, 255, 255), thickness=1, lineType=8)
return image_boxes
여기에 계산된 상자와 클래스 이름과 가능성의 값을 전달하고 여기에 4개의 값을 그리며 두 좌표는 모두 더 이상 정규화되지 않은 그림의 실제 픽셀 크기를 기준으로 한 값입니다.
그릴 왼쪽 위 오른쪽 아래 좌표 cv2를 꺼내고, name_bgr_dict에서 이 클래스가 어떤 색상에 해당해야 하는지 꺼내어 프레임을 그립니다.
그런 다음 상자의 왼쪽 위 모서리에서 시작 좌표로 텍스트를 만들고 오른쪽 아래 방향으로 작은 사각형을 그려서 다음과 같이 범주 이름과 확률을 채웁니다.
YOLODetector 클래스 파싱
def __init__(self,
model_path, class_name_list=None, mean_rgb=[122.67891434, 116.66876762, 104.00698793],
conf_thresh=0.1, prob_thresh=0.1, nms_thresh=0.5,
gpu_id=0):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
use_gpu = torch.cuda.is_available()
assert use_gpu, 'Current implementation does not support CPU mode. Enable CUDA.'
# Load YOLO model.
print("Loading YOLO model...")
self.yolo = resnet50()#这里就已经有了 随机的参数w权重
sd = torch.load(model_path)
self.yolo.load_state_dict(sd)#读取原来模型的权重
self.yolo.cuda()
print("Done loading!")
self.yolo.eval()
self.S = 7
self.B = 2
self.C = 20
self.class_name_list = class_name_list if (class_name_list is not None) else list(VOC_CLASS_BGR.keys())#给数据集里指定的list还是自己重新定义class list
assert len(self.class_name_list) == self.C
self.mean = np.array(mean_rgb, dtype=np.float32)
assert self.mean.shape == (3,)
self.conf_thresh = conf_thresh
self.prob_thresh = prob_thresh
self.nms_thresh = nms_thresh
self.to_tensor = transforms.ToTensor()
# Warm up. dummy_input 虚拟输入
dummy_input = Variable(torch.zeros((1, 3, 448, 448)))
dummy_input = dummy_input.cuda()
for i in range(3): #为了初始化权重？ 为什么 -预热操作的目的是让模型尽可能地填满加速器的缓存
self.yolo(dummy_input) #self.yolo.state_dict().get('conv1.weight')
yolo는 모델을 초기화하고 훈련된 model_path 위치의 가중치를 읽고 GPU에 넣습니다.
dummy_input을 사용하여 GPU를 예열하고 캐시를 먼저 점유합니다. 후속 감지 프로세스 중에 비디오 메모리, 메모리 또는 캐시가 버스트되지 않도록 방지합니다.
def detect(self, image_bgr, image_size=448):
""" Detect objects from given image.
Args:
image_bgr: (numpy array) input image in BGR ids_sorted, sized [h, w, 3].
image_size: (int) image width and height to which input image is resized.
Returns:
boxes_detected: (list of tuple) box corner list like [((x1, y1), (x2, y2))_obj1, ...]. Re-scaled for original input image size.
class_names_detected: (list of str) list of class name for each detected boxe.
probs_detected: (list of float) list of probability(=confidence x class_score) for each detected box.
"""
h, w, _ = image_bgr.shape
img = cv2.resize(image_bgr, dsize=(image_size, image_size), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # assuming the model is trained with RGB images.
img = (img - self.mean) / 255.0
img = self.to_tensor(img) # [image_size, image_size, 3] -> [3, image_size, image_size]
img = img[None, :, :, :] # [3, image_size, image_size] -> [1, 3, image_size, image_size]扩大维度 第一维是batch
img = Variable(img)
img = img.cuda()
with torch.no_grad():
pred_tensor = self.yolo(img)
pred_tensor = pred_tensor.cpu().data
pred_tensor = pred_tensor.squeeze(0) # squeeze batch dimension.
# Get detected boxes_detected, labels, confidences, class-scores.
boxes_normalized_all, class_labels_all, confidences_all, class_scores_all = self.decode(pred_tensor)
if boxes_normalized_all.size(0) == 0:
return [], [], [] # if no box found, return empty lists.
# Apply non maximum supression for boxes of each class.
boxes_normalized, class_labels, probs = [], [], []
for class_label in range(len(self.class_name_list)):
mask = (class_labels_all == class_label)
if torch.sum(mask) == 0:
continue # if no box found, skip that class.
# 找出所有同一类的 进行nms
boxes_normalized_masked = boxes_normalized_all[mask]
class_labels_maked = class_labels_all[mask]
confidences_masked = confidences_all[mask]
class_scores_masked = class_scores_all[mask]
ids = self.nms(boxes_normalized_masked, confidences_masked) #非极大抑制
boxes_normalized.append(boxes_normalized_masked[ids])
class_labels.append(class_labels_maked[ids])
probs.append(confidences_masked[ids] * class_scores_masked[ids])
boxes_normalized = torch.cat(boxes_normalized, 0)
class_labels = torch.cat(class_labels, 0)
probs = torch.cat(probs, 0)
# Postprocess for box, labels, probs.
boxes_detected, class_names_detected, probs_detected = [], [], []
for b in range(boxes_normalized.size(0)):
box_normalized = boxes_normalized[b]
class_label = class_labels[b]
prob = probs[b]
x1, x2 = w * box_normalized[0], w * box_normalized[2] # unnormalize x with image width. 图片真实坐标 从左上开始 0
y1, y2 = h * box_normalized[1], h * box_normalized[3] # unnormalize y with image height.
boxes_detected.append(((x1, y1), (x2, y2)))
class_label = int(class_label) # convert from LongTensor to int.
class_name = self.class_name_list[class_label]
class_names_detected.append(class_name)
prob = float(prob) # convert from Tensor to float.
probs_detected.append(prob)
return boxes_detected, class_names_detected, probs_detected
감지는 주요 기능을 감지합니다.
opencv에서 읽는 bgr이 rgb와 다르기 때문에 이것은 컴퓨터 비전의 역사에서 비롯된 것입니다. 나중에 cvt는 rgb로 변환됩니다.
voc2007의 평균값에서 그림의 uint8 값을 뺀 다음 정규화하여 255로 나눕니다. 평균값을 빼면 평균값은 0이 되고 /255의 분산은 1이 되어 정규분포를 따릅니다.
네트워크의 입력 크기는 (batch_size,3,448,448)이어야 합니다. 여기서 그림은 처음 3개의 차원만 가지고 있습니다. None을 사용하여 하나의 차원을 구성하고 이를 네트워크에 넣으면 감지를 예측할 때 예측값이 출력됩니다. At 동시에 이때 기울기를 계산할 필요가 없으며 사전에 no_grad를 설정할 필요가 없습니다. 그렇지 않으면 기계의 컴퓨팅 성능이 낭비됩니다. 그런 다음 출력 텐서를 가져오고 첫 번째 차원 배치 차원을 제거합니다.
디코더 디코드를 사용하여 네트워크 예측 출력의 텐서 pred_tensor 디코드를 분석합니다. 먼저 출력에 대해 이야기하겠습니다.
예측된 bbox 조건부 확률의 4개 그룹(객체를 포함하는 조건에서 범주의 확률을 가정)은 동일한 범주 중 가장 높습니다. txt 파일을 보면 사람, 즉 a라는 것을 알 수 있습니다. 사람
그런 다음 루프 조건으로 20개의 범주로 루프를 시작하고 마스크를 마스크하고 각 루프에서 예측된 범주가 이 루프의 범주에 있는지 확인합니다. 있는 경우 true를 제공하고 true에서 범주 내용을 계산합니다. 카테고리가 없습니다. 스킵하세요.
동일한 카테고리를 모두 찾아 nms를 수행하고 동일한 카테고리에 속하는 이들을 삭제하십시오. 모든 상자의 배치 및 조합의 iou가 너무 커서 중복으로 간주될 수도 있습니다. 이 상자를 지우고, 비교 값은 iou_threshold 임계값입니다.
nms가 완료되면 서로 겹치지 않는 상자의 인덱스를 얻습니다. 요약 목록에 추가합니다. 그런 다음 목록은 나중에 그림 프레임으로 반환되는 부분을 용이하게 하기 위해 행별로 텐서에 쌓입니다.
마지막 루프는 데이터를 출력 형식으로 다시 포맷합니다.
액자가 전달되는 곳은 4개의 값이 필요하고 2개의 좌표는 모두 더 이상 정규화되지 않은 사진의 실제 픽셀 크기를 기준으로 한 값이다. 여기서 해당 wh를 정규화하고 곱하여 실제 크기가 됩니다.
디코딩 부분
def decode(self, pred_tensor):
""" Decode tensor into box coordinates, class labels, and probs_detected.
Args:
pred_tensor: (tensor) tensor to decode sized [S, S, 5 x B + C], 5=(x, y, w, h, conf)
Returns:
boxes: (tensor) [[x1, y1, x2, y2]_obj1, ...]. Normalized from 0.0 to 1.0 w.r.t. image width/height, sized [n_boxes, 4].
labels: (tensor) class labels for each detected boxe, sized [n_boxes,].
confidences: (tensor) objectness confidences for each detected box, sized [n_boxes,].
class_scores: (tensor) scores for most likely class for each detected box, sized [n_boxes,].
"""
S, B, C = self.S, self.B, self.C
boxes, labels, confidences, class_scores = [], [], [], []
cell_size = 1.0 / float(S)
#每个网格的置信度
conf = pred_tensor[:, :, 4].unsqueeze(2) # [S, S, 1]
for b in range(1, B):
conf = torch.cat((conf, pred_tensor[:, :, 5*b + 4].unsqueeze(2)), 2) #[S,S,2]
conf_mask = conf > self.conf_thresh # [S, S, B]
# TBM, further optimization may be possible by replacing the following for-loops with tensor operations.
for i in range(S): # for x-dimension.
for j in range(S): # for y-dimension.
class_score, class_label = torch.max(pred_tensor[j, i, 5*B:], 0) #找[j,i]网格的最大分类值
for b in range(B): #遍历两预测bbox
conf = pred_tensor[j, i, 5*b + 4]
prob = conf * class_score
if float(prob) < self.prob_thresh: #低于阈值门限继续
continue
# Compute box corner (x1, y1, x2, y2) from tensor.
box = pred_tensor[j, i, 5*b : 5*b + 4]
x0y0_normalized = torch.FloatTensor([i, j]) * cell_size # 该网格的坐上角归一化坐标
xy_normalized = box[:2] * cell_size + x0y0_normalized # 从对cell归一化的中心点位置还原出来 现在是对图片大小归一化
wh_normalized = box[2:] # 归一化的宽高
box_xyxy = torch.FloatTensor(4) # [4,]随便初始4个
box_xyxy[:2] = xy_normalized - 0.5 * wh_normalized # 归一化左上X-》应该是左下角角位置(x1, y1).
box_xyxy[2:] = xy_normalized + 0.5 * wh_normalized # 归一化右下X-》应该是有右上角角位置(x2, y2).
# Append result to the lists.
boxes.append(box_xyxy)
labels.append(class_label)
confidences.append(conf)
class_scores.append(class_score)
if len(boxes) > 0:
boxes = torch.stack(boxes, 0) # [n_boxes, 4] list转张量
labels = torch.stack(labels, 0) # [n_boxes, ]
confidences = torch.stack(confidences, 0) # [n_boxes, ]
class_scores = torch.stack(class_scores, 0) # [n_boxes, ]
else:
# If no box found, return empty tensors.
boxes = torch.FloatTensor(0, 4)
labels = torch.LongTensor(0)
confidences = torch.FloatTensor(0)
class_scores = torch.FloatTensor(0)
return boxes, labels, confidences, class_scores예측된 텐서에 따라 상자 레이블(속해 있는 클래스의 첨자), 신뢰 범주 조건부 확률인 4개의 데이터 세트가 반환됩니다.
pred_tensor[:, :, 4]는 첫 번째 상자의 신뢰도를 제거한 다음 두 번째 상자의 신뢰도에 한 차원을 추가합니다.
코드의 conf_mask는 여기에서 사용되지 않으며 중요하지 않습니다.
가장 큰 범주에 해당하는 각 픽셀의 첨자와 확률 값을 얻기 위해 픽셀 단위로 반복합니다.
좌표 계산 순서는 박스의 4개 값(cx, cy, w, h)을 빼내고, 박스는 그리드 셀에서 오프셋되어 정규화된 실제 박스의 중심 값입니다. 셀의 xy 오프셋으로. 여기서는 그림을 기준으로 정규화된 상하좌우 좌표값으로 변환된다.
그 중 3단계는 2(B) 프레임을 순회하여 실제 확률보다 작은 임계값을 모두 빼는데 여기서 실제 확률을 conf * class_score(신뢰도 * 클래스 조건부 확률, 즉, 여기에 개체가 있을 확률 * 개체가 있는 경우 범주의 확률
확률이 요구 사항을 충족하면 이 상자에서 4개의 데이터를 꺼내서 4개의 요약 데이터 세트에 추가합니다.
4개의 목록 그룹이 텐서 형식으로 쌓이고 반환됩니다.
nms 부분
def nms(self, boxes, scores):
""" Apply non maximum supression.
Args:
Returns:
"""
threshold = self.nms_thresh
x1 = boxes[:, 0] # [n,]
y1 = boxes[:, 1] # [n,]
x2 = boxes[:, 2] # [n,]
y2 = boxes[:, 3] # [n,]
areas = (x2 - x1) * (y2 - y1) # [n,]
_, ids_sorted = scores.sort(0, descending=True) # [n,]
ids = []
while ids_sorted.numel() > 0:
# Assume `ids_sorted` size is [m,] in the beginning of this iter.
#最后剩下一个的时候detach 脱离出tensor
i = ids_sorted.item() if (ids_sorted.numel() == 1) else ids_sorted[0]
ids.append(i)
if ids_sorted.numel() == 1:
break # If only one box is left (i.e., no box to supress), break.
inter_x1 = x1[ids_sorted[1:]].clamp(min=x1[i]) # [m-1, ]
inter_y1 = y1[ids_sorted[1:]].clamp(min=y1[i]) # [m-1, ]
inter_x2 = x2[ids_sorted[1:]].clamp(max=x2[i]) # [m-1, ] 画图就懂了
inter_y2 = y2[ids_sorted[1:]].clamp(max=y2[i]) # [m-1, ]
inter_w = (inter_x2 - inter_x1).clamp(min=0) # [m-1, ]
inter_h = (inter_y2 - inter_y1).clamp(min=0) # [m-1, ]
inters = inter_w * inter_h # intersections b/w/ box `i` and other boxes, sized [m-1, ].
unions = areas[i] + areas[ids_sorted[1:]] - inters # unions b/w/ box `i` and other boxes, sized [m-1, ].
ious = inters / unions # [m-1, ]
# Remove boxes whose IoU is higher than the threshold.#(ious <= threshold).nonzero() 形状(2,1)
ids_keep = (ious <= threshold).nonzero().squeeze() # [m-1, ]. Because `nonzero()` adds extra dimension, squeeze it.
if ids_keep.numel() == 0:
break # If no box left, break.
ids_sorted = ids_sorted[ids_keep+1] # `+1` is needed because `ids_sorted[0] = i`.
return torch.LongTensor(ids)
nms는 비교적 하드코어인데 이해가 안되면 Li Mu의 13.4.Anchor frame-hands-on deep learning 2.0.0 문서를 읽어야 하지만 구현 방법이 약간 다릅니다.
상자의 그림을 기준으로 정규화된 왼쪽 위, 오른쪽 및 아래 좌표 값을 각각 꺼냅니다. 상자 면적 계산
점수(클래스의 조건부 확률)를 기준으로 최대에서 최소로 정렬합니다. inter의 x1 x2 y1 y2는 각각 비교할 상자의 왼쪽 하단과 오른쪽 상단에 해당하며,
점수는 순차적 벤치마크로 사용되며 매번 첫 번째(첨자 0)
x1을 확인하고 첨자 0인 x1과 비교합니다. 첨자 0인 x1보다 작으면 첨자 0인 x1로 자동 채워집니다.
y1을 확인하여 첨자 0인 y1과 비교합니다. 첨자 0인 y1보다 작으면 자동으로 첨자 0인 y1으로 채워집니다.
x2를 확인하고 첨자 0으로 x2와 비교하십시오. 첨자 0으로 x2보다 크면 자동으로 첨자 0으로 x2로 채워집니다.
y2를 확인하고 첨자 0인 y2와 비교합니다. 첨자 0인 y2보다 크면 자동으로 첨자 0인 y2로 채워집니다.
손으로 그림을 그리면 이해하기 쉽습니다.
그런 다음 둘의 w와 h를 계산합니다. 이때 두 상자는 겹치지 않고 교차하지 않습니다. w와 h 중 하나는 음수가 되고 0으로 채워지므로 iou는 0이어야 합니다. 어떤 교집합도 겹칠 수 없음을 의미하며 예약되고 0이 아닌 나머지 iou는 임계값과 비교되어 다음 그룹의 비교 대상으로 사용됩니다. 동시에 비교 벤치마크가 id에 입력되어 상자와 필터링된 비교 개체 간에 중복이 없음을 증명합니다.
이와 같이 나머지 다른 박스와 비교할 때마다 서로 겹치지 않거나 iou 임계값이 충분하지 않아 겹친 것으로 간주되는 모든 박스를 걸러내어 검출하여 처리하도록 한다.
총
import torch
from torch.autograd import Variable
import torchvision.transforms as transforms
import os
import cv2
import numpy as np
from resnet_yolo import resnet50
# VOC class names and BGR color.
VOC_CLASS_BGR = {
'aeroplane': (128, 0, 0),
'bicycle': (0, 128, 0),
'bird': (128, 128, 0),
'boat': (0, 0, 128),
'bottle': (128, 0, 128),
'bus': (0, 128, 128),
'car': (128, 128, 128),
'cat': (64, 0, 0),
'chair': (192, 0, 0),
'cow': (64, 128, 0),
'diningtable': (192, 128, 0),
'dog': (64, 0, 128),
'horse': (192, 0, 128),
'motorbike': (64, 128, 128),
'person': (192, 128, 128),
'pottedplant': (0, 64, 0),
'sheep': (128, 64, 0),
'sofa': (0, 192, 0),
'train': (128, 192, 0),
'tvmonitor': (0, 64, 128)
}
def visualize_boxes(image_bgr, boxes, class_names, probs, name_bgr_dict=None, line_thickness=2):
if name_bgr_dict is None:
name_bgr_dict = VOC_CLASS_BGR
image_boxes = image_bgr.copy()#分配到新内存中去
for box, class_name, prob in zip(boxes, class_names, probs):
# Draw box on the image.
left_top, right_bottom = box
left, top = int(left_top[0]), int(left_top[1])
right, bottom = int(right_bottom[0]), int(right_bottom[1])
bgr = name_bgr_dict[class_name]
cv2.rectangle(image_boxes, (left, top), (right, bottom), bgr, thickness=line_thickness)
# Draw text on the image.
text = '%s %.2f' % (class_name, prob)
size, baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=2)
text_w, text_h = size
x, y = left, top
x1y1 = (x, y)
x2y2 = (x + text_w + line_thickness, y + text_h + line_thickness + baseline)
cv2.rectangle(image_boxes, x1y1, x2y2, bgr, -1)
cv2.putText(image_boxes, text, (x + line_thickness, y + 2*baseline + line_thickness),
cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, color=(255, 255, 255), thickness=1, lineType=8)
return image_boxes
class YOLODetector:
def __init__(self,
model_path, class_name_list=None, mean_rgb=[122.67891434, 116.66876762, 104.00698793],
conf_thresh=0.1, prob_thresh=0.1, nms_thresh=0.5,
gpu_id=0):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
use_gpu = torch.cuda.is_available()
assert use_gpu, 'Current implementation does not support CPU mode. Enable CUDA.'
# Load YOLO model.
print("Loading YOLO model...")
self.yolo = resnet50()#这里就已经有了 随机的参数w权重
sd = torch.load(model_path)
self.yolo.load_state_dict(sd)#读取原来模型的权重
self.yolo.cuda()
print("Done loading!")
self.yolo.eval()
self.S = 7
self.B = 2
self.C = 20
self.class_name_list = class_name_list if (class_name_list is not None) else list(VOC_CLASS_BGR.keys())#给数据集里指定的list还是自己重新定义class list
assert len(self.class_name_list) == self.C
self.mean = np.array(mean_rgb, dtype=np.float32)
assert self.mean.shape == (3,)
self.conf_thresh = conf_thresh
self.prob_thresh = prob_thresh
self.nms_thresh = nms_thresh
self.to_tensor = transforms.ToTensor()
# Warm up. dummy_input 虚拟输入
dummy_input = Variable(torch.zeros((1, 3, 448, 448)))
dummy_input = dummy_input.cuda()
for i in range(3): #为了初始化权重？ 为什么 -预热操作的目的是让模型尽可能地填满加速器的缓存
self.yolo(dummy_input) #self.yolo.state_dict().get('conv1.weight')
def detect(self, image_bgr, image_size=448):
""" Detect objects from given image.
Args:
image_bgr: (numpy array) input image in BGR ids_sorted, sized [h, w, 3].
image_size: (int) image width and height to which input image is resized.
Returns:
boxes_detected: (list of tuple) box corner list like [((x1, y1), (x2, y2))_obj1, ...]. Re-scaled for original input image size.
class_names_detected: (list of str) list of class name for each detected boxe.
probs_detected: (list of float) list of probability(=confidence x class_score) for each detected box.
"""
h, w, _ = image_bgr.shape
img = cv2.resize(image_bgr, dsize=(image_size, image_size), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # assuming the model is trained with RGB images.
img = (img - self.mean) / 255.0
img = self.to_tensor(img) # [image_size, image_size, 3] -> [3, image_size, image_size]
img = img[None, :, :, :] # [3, image_size, image_size] -> [1, 3, image_size, image_size]扩大维度 第一维是batch
img = Variable(img)
img = img.cuda()
with torch.no_grad():
pred_tensor = self.yolo(img)
pred_tensor = pred_tensor.cpu().data
pred_tensor = pred_tensor.squeeze(0) # squeeze batch dimension.
# Get detected boxes_detected, labels, confidences, class-scores.
boxes_normalized_all, class_labels_all, confidences_all, class_scores_all = self.decode(pred_tensor)
if boxes_normalized_all.size(0) == 0:
return [], [], [] # if no box found, return empty lists.
# Apply non maximum supression for boxes of each class.
boxes_normalized, class_labels, probs = [], [], []
for class_label in range(len(self.class_name_list)):
mask = (class_labels_all == class_label)
if torch.sum(mask) == 0:
continue # if no box found, skip that class.
# 找出所有同一类的 进行nms
boxes_normalized_masked = boxes_normalized_all[mask]
class_labels_maked = class_labels_all[mask]
confidences_masked = confidences_all[mask]
class_scores_masked = class_scores_all[mask]
ids = self.nms(boxes_normalized_masked, confidences_masked) #非极大抑制
boxes_normalized.append(boxes_normalized_masked[ids])
class_labels.append(class_labels_maked[ids])
probs.append(confidences_masked[ids] * class_scores_masked[ids])
boxes_normalized = torch.cat(boxes_normalized, 0)
class_labels = torch.cat(class_labels, 0)
probs = torch.cat(probs, 0)
# Postprocess for box, labels, probs.
boxes_detected, class_names_detected, probs_detected = [], [], []
for b in range(boxes_normalized.size(0)):
box_normalized = boxes_normalized[b]
class_label = class_labels[b]
prob = probs[b]
x1, x2 = w * box_normalized[0], w * box_normalized[2] # unnormalize x with image width. 图片真实坐标 从左上开始 0
y1, y2 = h * box_normalized[1], h * box_normalized[3] # unnormalize y with image height.
boxes_detected.append(((x1, y1), (x2, y2)))
class_label = int(class_label) # convert from LongTensor to int.
class_name = self.class_name_list[class_label]
class_names_detected.append(class_name)
prob = float(prob) # convert from Tensor to float.
probs_detected.append(prob)
return boxes_detected, class_names_detected, probs_detected
def decode(self, pred_tensor):
""" Decode tensor into box coordinates, class labels, and probs_detected.
Args:
pred_tensor: (tensor) tensor to decode sized [S, S, 5 x B + C], 5=(x, y, w, h, conf)
Returns:
boxes: (tensor) [[x1, y1, x2, y2]_obj1, ...]. Normalized from 0.0 to 1.0 w.r.t. image width/height, sized [n_boxes, 4].
labels: (tensor) class labels for each detected boxe, sized [n_boxes,].
confidences: (tensor) objectness confidences for each detected box, sized [n_boxes,].
class_scores: (tensor) scores for most likely class for each detected box, sized [n_boxes,].
"""
S, B, C = self.S, self.B, self.C
boxes, labels, confidences, class_scores = [], [], [], []
cell_size = 1.0 / float(S)
#每个网格的置信度
conf = pred_tensor[:, :, 4].unsqueeze(2) # [S, S, 1]
for b in range(1, B):
conf = torch.cat((conf, pred_tensor[:, :, 5*b + 4].unsqueeze(2)), 2) #[S,S,2]
conf_mask = conf > self.conf_thresh # [S, S, B]
# TBM, further optimization may be possible by replacing the following for-loops with tensor operations.
for i in range(S): # for x-dimension.
for j in range(S): # for y-dimension.
class_score, class_label = torch.max(pred_tensor[j, i, 5*B:], 0) #找[j,i]网格的最大分类值
for b in range(B): #遍历两预测bbox
conf = pred_tensor[j, i, 5*b + 4]
prob = conf * class_score
if float(prob) < self.prob_thresh: #低于阈值门限继续
continue
# Compute box corner (x1, y1, x2, y2) from tensor.
box = pred_tensor[j, i, 5*b : 5*b + 4]
x0y0_normalized = torch.FloatTensor([i, j]) * cell_size # 该网格的坐上角归一化坐标
xy_normalized = box[:2] * cell_size + x0y0_normalized # 从对cell归一化的中心点位置还原出来 现在是对图片大小归一化
wh_normalized = box[2:] # 归一化的宽高
box_xyxy = torch.FloatTensor(4) # [4,]随便初始4个
box_xyxy[:2] = xy_normalized - 0.5 * wh_normalized # 归一化左上X-》应该是左下角角位置(x1, y1).
box_xyxy[2:] = xy_normalized + 0.5 * wh_normalized # 归一化右下X-》应该是有右上角角位置(x2, y2).
# Append result to the lists.
boxes.append(box_xyxy)
labels.append(class_label)
confidences.append(conf)
class_scores.append(class_score)
if len(boxes) > 0:
boxes = torch.stack(boxes, 0) # [n_boxes, 4] list转张量
labels = torch.stack(labels, 0) # [n_boxes, ]
confidences = torch.stack(confidences, 0) # [n_boxes, ]
class_scores = torch.stack(class_scores, 0) # [n_boxes, ]
else:
# If no box found, return empty tensors.
boxes = torch.FloatTensor(0, 4)
labels = torch.LongTensor(0)
confidences = torch.FloatTensor(0)
class_scores = torch.FloatTensor(0)
return boxes, labels, confidences, class_scores
def nms(self, boxes, scores):
""" Apply non maximum supression.
Args:
Returns:
"""
threshold = self.nms_thresh
x1 = boxes[:, 0] # [n,]
y1 = boxes[:, 1] # [n,]
x2 = boxes[:, 2] # [n,]
y2 = boxes[:, 3] # [n,]
areas = (x2 - x1) * (y2 - y1) # [n,]
_, ids_sorted = scores.sort(0, descending=True) # [n,]
ids = []
while ids_sorted.numel() > 0:
# Assume `ids_sorted` size is [m,] in the beginning of this iter.
#最后剩下一个的时候detach 脱离出tensor
i = ids_sorted.item() if (ids_sorted.numel() == 1) else ids_sorted[0]
ids.append(i)
if ids_sorted.numel() == 1:
break # If only one box is left (i.e., no box to supress), break.
inter_x1 = x1[ids_sorted[1:]].clamp(min=x1[i]) # [m-1, ]
inter_y1 = y1[ids_sorted[1:]].clamp(min=y1[i]) # [m-1, ]
inter_x2 = x2[ids_sorted[1:]].clamp(max=x2[i]) # [m-1, ] 画图就懂了
inter_y2 = y2[ids_sorted[1:]].clamp(max=y2[i]) # [m-1, ]
inter_w = (inter_x2 - inter_x1).clamp(min=0) # [m-1, ]
inter_h = (inter_y2 - inter_y1).clamp(min=0) # [m-1, ]
inters = inter_w * inter_h # intersections b/w/ box `i` and other boxes, sized [m-1, ].
unions = areas[i] + areas[ids_sorted[1:]] - inters # unions b/w/ box `i` and other boxes, sized [m-1, ].
ious = inters / unions # [m-1, ]
# Remove boxes whose IoU is higher than the threshold.#(ious <= threshold).nonzero() 形状(2,1)
ids_keep = (ious <= threshold).nonzero().squeeze() # [m-1, ]. Because `nonzero()` adds extra dimension, squeeze it.
if ids_keep.numel() == 0:
break # If no box left, break.
ids_sorted = ids_sorted[ids_keep+1] # `+1` is needed because `ids_sorted[0] = i`.
return torch.LongTensor(ids)
if __name__ == '__main__':
# Paths to input/output images.
image_path = '000369.jpg'
out_path = 'result.png'
# Path to the yolo weight.
model_path = 'weights/model_best.pth'
# GPU device on which yolo is loaded.
gpu_id = 0
# Load model.
yolo = YOLODetector(model_path, gpu_id=gpu_id, conf_thresh=0.15, prob_thresh=0.45, nms_thresh=0.35)
# Load image.
image = cv2.imread(image_path)#某些老的图像处理软件使用的是 BGR 格式，因此 OpenCV 采用 BGR 格式可以与这些软件兼容。
# Detect objects.
boxes, class_names, probs = yolo.detect(image)
# Visualize.
image_boxes = visualize_boxes(image, boxes, class_names, probs)
# Output detection result as an image.
cv2.imwrite(out_path, image_boxes)