- import numpy as np
- #k是卷积核假设n*n,feature是特征图假设是w*h
- def convolution(feature,k):
- w,h=feature.shape
- n=len(k)
- new_feature=[]
- for i in range(w-n):
- line=[]
- for j in range(h-n):
- a=feature[i:i+n,j:j+n]
- #np.multiply()求两个矩阵的内积
- line.append(np.sum(np.multiply(k,a)))
- new_feature.append(line)
- return np.array(new_feature)
- import torch
- #坐标原点在左上角,水平x,竖直y
- #计算一个框与一堆框的iou
- def iou(box,boxes):
- #先计算box的面积[x1,y1,x2,y2]
- box_area=(box[2]-box[0])*(box[3]-box[1])
- boxes_area=(boxes[:,2]-boxes[:,0])*(boxes[:,3]-boxes[:,1])
- #求交集,左上角最大,右下角最小
- xx1=torch.maximum(box[0],boxes[:,0])
- yy1=torch.maximum(box[1],boxes[:,1])
- xx2=torch.minimum(box[2],boxes[:,2])
- yy2=torch.minimum(box[3],boxes[:,3])
- w,h=torch.maximum(torch.Tensor([0]),xx2-xx1),torch.maximum(torch.Tensor([0]),yy2-yy1)
- over_area=w*h
- return over_area/(box_area+boxes_area-over_area)
-
- def NMS(boxes,thresh=0.3):
- #根据boxes的置信度进行排序,假设置信度在第一列
- #argsort将numpy数组进行排序,返回索引,descending表示从大到小返回
- new_boxes=boxes[boxes[:,0].argsort(descending=True)]
- keep_boxes=[]
- while len(new_boxes)>0:
- max_box=new_boxes[0]
- keep_boxes.append(max_box)
- if len(new_boxes)>1:#除了评分最高的外,还有其他框
- other_boxes=new_boxes[1:]
- #torch.where返回满足条件的索引
- new_boxes=other_boxes[torch.where(iou(max_box[1:],other_boxes[:,1:])
- else:
- break
- #将一个列表里面的张量数据进行扩维拼接,本来1行5列,现在n行5列
- return torch.stack(keep_boxes)
- boxes=torch.Tensor([[0.8,0,0,2,2],[0.6,1,1,4,4],[0.5,-1,-1,2,2],[0.75,-0.5,-0.5,2.5,2.5]])
- print(NMS(boxes,0.3))
numpy版本
- def nms(boxes, scores, threshold):
- # 根据得分降序排序
- indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
-
- # 初始化选中的边界框列表
- selected_boxes = []
-
- while indices:
- # 选择得分最高的边界框
- i = indices[0]
- selected_boxes.append(boxes[i])
-
- # 计算当前边界框与其他边界框的IoU
- ious = [calculate_iou(boxes[i], boxes[j]) for j in indices[1:]]
-
- # 仅保留IoU低于阈值的边界框
- indices = [indices[j + 1] for j, iou in enumerate(ious) if iou < threshold]
-
- return selected_boxes
-
-
- def calculate_iou(box1, box2):
- # 计算两个边界框的相交区域
- x1 = max(box1[0], box2[0])
- y1 = max(box1[1], box2[1])
- x2 = min(box1[2], box2[2])
- y2 = min(box1[3], box2[3])
-
- intersection = max(0, x2 - x1) * max(0, y2 - y1)
-
- # 计算两个边界框的联合区域
- area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
- area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
-
- union = area1 + area2 - intersection
-
- # 计算IoU(交并比)
- iou = intersection / union
-
- return iou