非
极
大
值
抑
制
非极大值抑制
非 极 大 值 抑 制
import torch
def bbox_iou ( box1, box2, x1y1x2y2= True ) :
"""
计算IOU
"""
if not x1y1x2y2:
b1_x1, b1_x2 = box1[ : , 0 ] - box1[ : , 2 ] / 2 , box1[ : , 0 ] + box1[ : , 2 ] / 2
b1_y1, b1_y2 = box1[ : , 1 ] - box1[ : , 3 ] / 2 , box1[ : , 1 ] + box1[ : , 3 ] / 2
b2_x1, b2_x2 = box2[ : , 0 ] - box2[ : , 2 ] / 2 , box2[ : , 0 ] + box2[ : , 2 ] / 2
b2_y1, b2_y2 = box2[ : , 1 ] - box2[ : , 3 ] / 2 , box2[ : , 1 ] + box2[ : , 3 ] / 2
else :
b1_x1, b1_y1, b1_x2, b1_y2 = box1[ : , 0 ] , box1[ : , 1 ] , box1[ : , 2 ] , box1[ : , 3 ]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[ : , 0 ] , box2[ : , 1 ] , box2[ : , 2 ] , box2[ : , 3 ]
inter_rect_x1 = torch. max ( b1_x1, b2_x1)
inter_rect_y1 = torch. max ( b1_y1, b2_y1)
inter_rect_x2 = torch. min ( b1_x2, b2_x2)
inter_rect_y2 = torch. min ( b1_y2, b2_y2)
inter_area = torch. clamp( inter_rect_x2 - inter_rect_x1 + 1e - 3 , min = 0 ) * \
torch. clamp( inter_rect_y2 - inter_rect_y1 + 1e - 3 , min = 0 )
b1_area = ( b1_x2 - b1_x1 + 1e - 3 ) * ( b1_y2 - b1_y1 + 1e - 3 )
b2_area = ( b2_x2 - b2_x1 + 1e - 3 ) * ( b2_y2 - b2_y1 + 1e - 3 )
iou = inter_area / ( b1_area + b2_area - inter_area + 1e - 16 )
return iou
def non_max_suppression ( prediction, num_classes, conf_thres= 0.5 , nms_thres= 0.4 ) :
box_corner = prediction. new( prediction. shape)
box_corner[ : , : , 0 ] = prediction[ : , : , 0 ] - prediction[ : , : , 2 ] / 2
box_corner[ : , : , 1 ] = prediction[ : , : , 1 ] - prediction[ : , : , 3 ] / 2
box_corner[ : , : , 2 ] = prediction[ : , : , 0 ] + prediction[ : , : , 2 ] / 2
box_corner[ : , : , 3 ] = prediction[ : , : , 1 ] + prediction[ : , : , 3 ] / 2
prediction[ : , : , : 4 ] = box_corner[ : , : , : 4 ]
output = [ None for _ in range ( len ( prediction) ) ]
for image_i, image_pred in enumerate ( prediction) :
conf_mask = ( image_pred[ : , 4 ] >= conf_thres) . squeeze( )
image_pred = image_pred[ conf_mask]
if not image_pred. size( 0 ) :
continue
class_conf, class_pred = torch. max ( image_pred[ : , 5 : 5 + num_classes] , 1 , keepdim= True )
detections = torch. cat( ( image_pred[ : , : 5 ] , class_conf. float ( ) , class_pred. float ( ) ) , 1 )
with torch. no_grad( ) :
unique_labels = detections[ : , - 1 ] . cpu( ) . unique( )
if prediction. is_cuda:
unique_labels = unique_labels. cuda( )
for c in unique_labels:
detections_class = detections[ detections[ : , - 1 ] == c]
_, conf_sort_index = torch. sort( detections_class[ : , 4 ] , descending= True )
detections_class = detections_class[ conf_sort_index]
max_detections = [ ]
while detections_class. size( 0 ) :
max_detections. append( detections_class[ 0 ] . unsqueeze( 0 ) )
if len ( detections_class) == 1 :
break
ious = bbox_iou( max_detections[ - 1 ] , detections_class[ 1 : ] )
detections_class = detections_class[ 1 : ] [ ious < nms_thres]
max_detections = torch. cat( max_detections) . data
output[ image_i] = max_detections if output[ image_i] is None else torch. cat(
( output[ image_i] , max_detections) )
return output
if __name__ == '__main__' :
from utils import get_anchors
from yolo4 import YoloBody
from yolo_layer import YoloLayer
rgb = torch. randn( 1 , 3 , 608 , 608 )
model = YoloBody( 3 , 80 )
out = model( rgb)
anchors = get_anchors( 'yolo_anchors_coco.txt' )
yolo_decodes = [ ]
anchor_masks = [ [ 0 , 1 , 2 ] , [ 3 , 4 , 5 ] , [ 6 , 7 , 8 ] ]
for i in range ( 3 ) :
head = YoloLayer( ( 608 , 608 , 3 ) , anchor_masks, 80 ,
anchors, len ( anchors) // 2 ) . eval ( )
yolo_decodes. append( head)
output_list = [ ]
for i in range ( 3 ) :
output_list. append( yolo_decodes[ i] ( out[ i] ) )
output = torch. cat( output_list, 1 )
batch_detections = non_max_suppression( output, 80 )
boxes = [ box. cpu( ) . numpy( ) for box in batch_detections]
print ( boxes[ 0 ] )
def get_anchors ( yolo_anchors_path) :
import os
anchors_path = os. path. expanduser( yolo_anchors_path)
with open ( anchors_path) as f:
anchors = f. readline( )
anchors = [ float ( x) for x in anchors. split( ',' ) ]
return anchors
import torch
import torch. nn as nn
from collections import OrderedDict
from torch. autograd import Variable
from CSPDarknet import *
import cv2
def conv2d ( filter_in, filter_out, kernel_size, stride= 1 ) :
pad = ( kernel_size - 1 ) // 2 if kernel_size else 0
return nn. Sequential( OrderedDict( [
( "conv" , nn. Conv2d( filter_in, filter_out, kernel_size= kernel_size, stride= stride, padding= pad, bias= False ) ) ,
( "bn" , nn. BatchNorm2d( filter_out) ) ,
( "relu" , nn. LeakyReLU( 0.1 ) ) ,
] ) )
class SpatialPyramidPooling ( nn. Module) :
def __init__ ( self, pool_sizes= [ 5 , 9 , 13 ] ) :
super ( SpatialPyramidPooling, self) . __init__( )
self. maxpools = nn. ModuleList( [ nn. MaxPool2d( pool_size, 1 , pool_size// 2 ) for pool_size in pool_sizes] )
def forward ( self, x) :
features = [ maxpool( x) for maxpool in self. maxpools[ : : - 1 ] ]
features = torch. cat( features + [ x] , dim= 1 )
return features
class Upsample ( nn. Module) :
def __init__ ( self, in_channels, out_channels) :
super ( Upsample, self) . __init__( )
self. upsample = nn. Sequential(
conv2d( in_channels, out_channels, 1 ) ,
nn. Upsample( scale_factor= 2 , mode= 'nearest' )
)
def forward ( self, x, ) :
x = self. upsample( x)
return x
def make_three_conv ( filters_list, in_filters) :
m = nn. Sequential(
conv2d( in_filters, filters_list[ 0 ] , 1 ) ,
conv2d( filters_list[ 0 ] , filters_list[ 1 ] , 3 ) ,
conv2d( filters_list[ 1 ] , filters_list[ 0 ] , 1 ) ,
)
return m
def make_five_conv ( filters_list, in_filters) :
m = nn. Sequential(
conv2d( in_filters, filters_list[ 0 ] , 1 ) ,
conv2d( filters_list[ 0 ] , filters_list[ 1 ] , 3 ) ,
conv2d( filters_list[ 1 ] , filters_list[ 0 ] , 1 ) ,
conv2d( filters_list[ 0 ] , filters_list[ 1 ] , 3 ) ,
conv2d( filters_list[ 1 ] , filters_list[ 0 ] , 1 ) ,
)
return m
def yolo_head ( filters_list, in_filters) :
m = nn. Sequential(
conv2d( in_filters, filters_list[ 0 ] , 3 ) ,
nn. Conv2d( filters_list[ 0 ] , filters_list[ 1 ] , 1 ) ,
)
return m
class YoloBody ( nn. Module) :
def __init__ ( self, num_anchors, num_classes) :
super ( YoloBody, self) . __init__( )
self. backbone = darknet53( None )
self. conv1 = make_three_conv( [ 512 , 1024 ] , 1024 )
self. SPP = SpatialPyramidPooling( )
self. conv2 = make_three_conv( [ 512 , 1024 ] , 2048 )
self. upsample1 = Upsample( 512 , 256 )
self. conv_for_P4 = conv2d( 512 , 256 , 1 )
self. make_five_conv1 = make_five_conv( [ 256 , 512 ] , 512 )
self. upsample2 = Upsample( 256 , 128 )
self. conv_for_P3 = conv2d( 256 , 128 , 1 )
self. make_five_conv2 = make_five_conv( [ 128 , 256 ] , 256 )
final_out_filter2 = num_anchors * ( 5 + num_classes)
self. yolo_head3 = yolo_head( [ 256 , final_out_filter2] , 128 )
self. down_sample1 = conv2d( 128 , 256 , 3 , stride= 2 )
self. make_five_conv3 = make_five_conv( [ 256 , 512 ] , 512 )
final_out_filter1 = num_anchors * ( 5 + num_classes)
self. yolo_head2 = yolo_head( [ 512 , final_out_filter1] , 256 )
self. down_sample2 = conv2d( 256 , 512 , 3 , stride= 2 )
self. make_five_conv4 = make_five_conv( [ 512 , 1024 ] , 1024 )
final_out_filter0 = num_anchors * ( 5 + num_classes)
self. yolo_head1 = yolo_head( [ 1024 , final_out_filter0] , 512 )
def forward ( self, x) :
x2, x1, x0 = self. backbone( x)
P5 = self. conv1( x0)
P5 = self. SPP( P5)
P5 = self. conv2( P5)
P5_upsample = self. upsample1( P5)
P4 = self. conv_for_P4( x1)
P4 = torch. cat( [ P4, P5_upsample] , axis= 1 )
P4 = self. make_five_conv1( P4)
P4_upsample = self. upsample2( P4)
P3 = self. conv_for_P3( x2)
P3 = torch. cat( [ P3, P4_upsample] , axis= 1 )
P3 = self. make_five_conv2( P3)
P3_downsample = self. down_sample1( P3)
P4 = torch. cat( [ P3_downsample, P4] , axis= 1 )
P4 = self. make_five_conv3( P4)
P4_downsample = self. down_sample2( P4)
P5 = torch. cat( [ P4_downsample, P5] , axis= 1 )
P5 = self. make_five_conv4( P5)
out2 = self. yolo_head3( P3)
out1 = self. yolo_head2( P4)
out0 = self. yolo_head1( P5)
return out0, out1, out2
if __name__ == '__main__' :
rgb = torch. randn( 1 , 3 , 608 , 608 )
model = YoloBody( 3 , 80 )
load_model_pth( model, 'pth/yolo4_weights_my.pth' )
out = model( rgb)
print ( '-----' * 5 )
print ( out[ 0 ] . shape)
print ( '-----' * 5 )
print ( out[ 1 ] . shape)
print ( '-----' * 5 )
print ( out[ 2 ] . shape)
print ( '-----' * 5 )
import torch. nn as nn
import torch. nn. functional as F
import torch
import numpy as np
def yolo_decode ( output, num_classes, anchors, num_anchors, scale_x_y) :
device = None
cuda_check = output. is_cuda
if cuda_check:
device = output. get_device( )
n_ch = 4 + 1 + num_classes
A = num_anchors
B = output. size( 0 )
H = output. size( 2 )
W = output. size( 3 )
output = output. view( B, A, n_ch, H, W) . permute( 0 , 1 , 3 , 4 , 2 ) . contiguous( )
bx, by = output[ . . . , 0 ] , output[ . . . , 1 ]
bw, bh = output[ . . . , 2 ] , output[ . . . , 3 ]
det_confs = output[ . . . , 4 ]
cls_confs = output[ . . . , 5 : ]
bx = torch. sigmoid( bx)
by = torch. sigmoid( by)
bw = torch. exp( bw) * scale_x_y - 0.5 * ( scale_x_y- 1 )
bh = torch. exp( bh) * scale_x_y - 0.5 * ( scale_x_y- 1 )
det_confs = torch. sigmoid( det_confs)
cls_confs = torch. sigmoid( cls_confs)
grid_x = torch. arange( W, dtype= torch. float ) . repeat( 1 , 3 , W, 1 ) . to( device)
grid_y = torch. arange( H, dtype= torch. float ) . repeat( 1 , 3 , H, 1 ) . permute( 0 , 1 , 3 , 2 ) . to( device)
bx += grid_x
by += grid_y
for i in range ( num_anchors) :
bw[ : , i, : , : ] *= anchors[ i* 2 ]
bh[ : , i, : , : ] *= anchors[ i* 2 + 1 ]
bx = ( bx / W) . unsqueeze( - 1 )
by = ( by / H) . unsqueeze( - 1 )
bw = ( bw / W) . unsqueeze( - 1 )
bh = ( bh / H) . unsqueeze( - 1 )
boxes = torch. cat( ( bx, by, bw, bh) , dim= - 1 ) . reshape( B, A * H * W, 4 )
det_confs = det_confs. unsqueeze( - 1 ) . reshape( B, A* H* W, 1 )
cls_confs = cls_confs. reshape( B, A* H* W, num_classes)
outputs = torch. cat( [ boxes, det_confs, cls_confs] , dim= - 1 )
return outputs
class YoloLayer ( nn. Module) :
''' Yolo layer
model_out: while inference,is post-processing inside or outside the model
true:outside
'''
def __init__ ( self, img_size, anchor_masks= [ ] , num_classes= 80 , anchors= [ ] , num_anchors= 9 , scale_x_y= 1 ) :
super ( YoloLayer, self) . __init__( )
self. anchor_masks = anchor_masks
self. num_classes = num_classes
if type ( anchors) == np. ndarray:
self. anchors = anchors. tolist( )
else :
self. anchors = anchors
print ( self. anchors)
print ( type ( self. anchors) )
self. num_anchors = num_anchors
self. anchor_step = len ( self. anchors) // num_anchors
print ( self. anchor_step)
self. scale_x_y = scale_x_y
self. feature_length = [ img_size[ 0 ] // 8 , img_size[ 0 ] // 16 , img_size[ 0 ] // 32 ]
self. img_size = img_size
def forward ( self, output) :
if self. training:
return output
in_w = output. size( 3 )
anchor_index = self. anchor_masks[ self. feature_length. index( in_w) ]
stride_w = self. img_size[ 0 ] / in_w
masked_anchors = [ ]
for m in anchor_index:
masked_anchors += self. anchors[ m * self. anchor_step: ( m + 1 ) * self. anchor_step]
self. masked_anchors = [ anchor / stride_w for anchor in masked_anchors]
data = yolo_decode( output, self. num_classes, self. masked_anchors, len ( anchor_index) , scale_x_y= self. scale_x_y)
return data
if __name__ == '__main__' :
from utils import get_anchors
from yolo4 import YoloBody
rgb = torch. randn( 1 , 3 , 608 , 608 )
model = YoloBody( 3 , 80 )
out = model( rgb)
anchors = get_anchors( 'yolo_anchors_coco.txt' )
yolo_decodes = [ ]
anchor_masks = [ [ 0 , 1 , 2 ] , [ 3 , 4 , 5 ] , [ 6 , 7 , 8 ] ]
for i in range ( 3 ) :
head = YoloLayer( ( 608 , 608 , 3 ) , anchor_masks, 80 ,
anchors, len ( anchors) // 2 ) . eval ( )
yolo_decodes. append( head)
output_list = [ ]
for i in range ( 3 ) :
output_list. append( yolo_decodes[ i] ( out[ i] ) )
output = torch. cat( output_list, 1 )
print ( output. shape)