yolov3后处理，包括网络输出、阈值过滤、多类NMS

原创

已于 2022-06-15 17:01:23 修改 · 1.4k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习

于 2022-06-15 16:52:35 首次发布

pytorch版本：1.7.1+cu101

from net.yolov3 import Yolo
from PIL import Image, ImageDraw
import numpy as np
import torch, copy, time

def sigmoid(x):
	temporary = 1 + torch.exp(-x)
	return 1.0 / temporary

def get_boxes(output, anchors):
	h = output.size(2)
	w = output.size(3)
	output = output.view(3, 85, h, w).permute(0,2,3,1)
	tc = torch.sigmoid(output[..., 4])    # 3*h*w
	cl = torch.sigmoid(output[..., 5:])   # 3*h*w*80
	clv, cli = torch.max(cl,-1)
	mask = tc * clv > 0.9
	# print(torch.where(mask))
	cli = cli[mask].unsqueeze(-1)
	# 3*h*w
	tx = torch.sigmoid(output[..., 0][mask])
	ty = torch.sigmoid(output[..., 1][mask])
	tw = torch.exp(output[..., 2][mask])
	th = torch.exp(output[..., 3][mask])
	# grid
	FloatTensor = torch.cuda.FloatTensor if tx.is_cuda else torch.FloatTensor
	grid_x, grid_y = torch.meshgrid(torch.linspace(0,w-1,w), torch.linspace(0,h-1,h))
	grid_x = grid_x.repeat(3,1,1)[mask].type(FloatTensor)
	grid_y = grid_y.repeat(3,1,1)[mask].type(FloatTensor)
	tx = ((tx+grid_y) / w).unsqueeze(-1)
	ty = ((ty+grid_x) / h).unsqueeze(-1)
	# anchor
	aw = torch.Tensor(anchors[0::2]).view(3,1).repeat(1,h*w).view(3,h,w)[mask].type(FloatTensor)
	ah = torch.Tensor(anchors[1::2]).view(3,1