295 lines
12 KiB
Plaintext
295 lines
12 KiB
Plaintext
import sys
|
||
from pathlib import Path
|
||
import math
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
|
||
FILE = Path(__file__).absolute()
|
||
#sys.path.append(FILE.parents[0].as_posix()) # add yolov5/ to path
|
||
'''
|
||
修改说明：
|
||
1、pars中增加了recScale参数。船舶判断是否悬挂国旗时，需将船舶检测框乘以扩大系数imgScale后，与国旗中心点坐标比较。
|
||
pars={'imgSize':(imgwidth,imgheight),'wRation':1/6.0,'hRation':1/6.0,'smallId':0,'bigId':3,'newId':4,'recScale':1.2}
|
||
2、增加expand_rect(preds_boat, recScale, imgSize)函数，在图像范围内，将矩形框扩大recScale倍数。
|
||
3、增加或修改以下两行：
|
||
preds_boat_flag_expand=expand_rect(preds_boat_flag[i],pars['recScale'],pars['imgSize']) #新增！
|
||
if point_in_rectangle(preds_flag,preds_boat_flag_expand)>=1: #新增后修改！
|
||
'''
|
||
|
||
def channel2_post_process(predsList,pars):
|
||
|
||
#pars={'imgSize':(imgwidth,imgheight),'wRation':1/6.0,'hRation':1/6.0,'smallId':0,'bigId':3,'newId':4,'recScale':1.2}
|
||
'''
|
||
后处理程序，将检测结果中未悬挂国旗的船只，其类别改为4，即'unflagged_ship'
|
||
最终类别汇总如下，
|
||
['flag', 'buoy', 'shipname', 'ship','unflagged_ship']=[0,1,2,3,4]
|
||
|
||
输入：
|
||
preds 一张图像的检测结果，为嵌套列表，tensor，包括x_y_x_y_conf_class
|
||
imgwidth,imgheight 图像的原始宽度及长度
|
||
输出：检测结果（将其中未悬挂国旗的显示）
|
||
'''
|
||
|
||
preds = torch.tensor(predsList[0])
|
||
preds=preds.tolist()
|
||
print('---line36:',preds)
|
||
preds = filter_detection_results(preds,pars)
|
||
|
||
|
||
preds=[[*sublist[:-1], int(sublist[-1])] for sublist in preds] #类别从浮点型转为整型
|
||
|
||
#设置空的列表
|
||
output_detection=[] #存放往接口传的类别
|
||
#1、判断类别中哪些有船？取出船检测结果，并取出国旗检测结果。
|
||
# output_detection.append[] 这里将船和国旗以外的类别加进去
|
||
preds_boat=[]
|
||
preds_flag=[]
|
||
|
||
#jcq： 增加封仓
|
||
preds_uncover = []
|
||
|
||
for i in range(len(preds)):
|
||
if preds[i][5]==pars['bigId']: #识别为船
|
||
preds_boat.append(preds[i])
|
||
elif preds[i][5]==pars['smallId']: #识别为国旗
|
||
preds_flag.append(preds[i])
|
||
output_detection.append(preds[i])
|
||
|
||
#jcq:
|
||
elif preds[i][5]==pars['uncoverId']: #识封仓
|
||
preds_uncover.append(preds[i])
|
||
output_detection.append(preds[i])
|
||
|
||
else:
|
||
output_detection.append(preds[i])
|
||
# pass
|
||
|
||
# return output_detection
|
||
|
||
#jcq: 自动检测船只 - 0311
|
||
#2、船尺寸与图像比较，其中长或宽有一个维度超过图像宽高平均值的1/3，启动国旗检测
|
||
#①if 判断：判断超过1/3的，则取出这些大船，进一步判断是否悬挂国旗
|
||
#不超过1/3的，则output_detection.append[]
|
||
|
||
preds_boat_flag=[]
|
||
for i in range(len(preds_boat)):
|
||
length_bbx,width_bbx=get_rectangle_dimensions(preds_boat[i])
|
||
length_bbx, width_bbx=int(length_bbx),int(width_bbx)
|
||
if length_bbx>(pars['imgSize'][0]+pars['imgSize'][1])* pars['hRation'] or width_bbx>(pars['imgSize'][0]+pars['imgSize'][1])*pars['wRation']:
|
||
preds_boat_flag.append(preds_boat[i])
|
||
else:
|
||
output_detection.append(preds_boat[i])
|
||
|
||
|
||
#②将大船的框与国旗检测结果的中心点坐标做比较。
|
||
#若没有一个在，则输出此船未悬挂国旗（船舶类别名称改完未悬挂国旗就行，即将0、1、2、3中的0替换为4的类别）
|
||
# 未悬挂国旗的，则output_detection.append[xyxy_4_conf]
|
||
#若有国旗在，则不输出是否悬挂国旗，则output_detection.append[xyxy_0_conf]
|
||
|
||
#jcq : 判断是否有国旗 - 具体判断方法使用交叉比
|
||
for i in range(len(preds_boat_flag)):
|
||
preds_boat_flag_expand=expand_rect(preds_boat_flag[i],pars['recScale'],pars['imgSize']) #新增！
|
||
if point_in_rectangle(preds_flag,preds_boat_flag_expand)>=1: #新增后修改！
|
||
output_detection.append(preds_boat_flag[i])
|
||
else:
|
||
temp_preds_boat_flag=preds_boat_flag[i]
|
||
temp_preds_boat_flag[5]=pars['newId'] #将类别标签改为4，即为未悬挂国旗的船只
|
||
output_detection.append(temp_preds_boat_flag)
|
||
|
||
|
||
#jcq: 将大船的框与封仓检测结果的中心点坐标做比较。 -- 该步可以暂时不做，因为船只比封仓更容易检测出结果
|
||
# 若未封仓，则输出此船未封仓（船舶类别名称改船只封仓，即将0、1、2、3中的3替换为未封仓的类别）
|
||
# 未封仓的，则output_detection.append[xyxy_4_conf]
|
||
# 若有国旗在，则不输出是否悬挂国旗，则output_detection.append[xyxy_0_conf]
|
||
|
||
# for i in range(len(preds_boat_flag)):
|
||
# preds_boat_flag_expand=expand_rect(preds_boat_flag[i],pars['recScale'],pars['imgSize']) #新增！
|
||
# if point_in_rectangle(preds_flag,preds_boat_flag_expand)>=1: #新增后修改！
|
||
# output_detection.append(preds_boat_flag[i])
|
||
# else:
|
||
# temp_preds_boat_flag=preds_boat_flag[i]
|
||
# temp_preds_boat_flag[5]=pars['newId'] #将类别标签改为4，即为未悬挂国旗的船只
|
||
# output_detection.append(temp_preds_boat_flag)
|
||
|
||
|
||
|
||
|
||
return output_detection
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
def center_coordinate(boundbxs):
|
||
'''
|
||
根据检测矩形框，得到其矩形长度和宽度
|
||
输入：两个对角坐标xyxy
|
||
输出：矩形框重点坐标xy
|
||
'''
|
||
boundbxs_x1 = boundbxs[0]
|
||
boundbxs_y1 = boundbxs[1]
|
||
boundbxs_x2 = boundbxs[2]
|
||
boundbxs_y2 = boundbxs[3]
|
||
center_x = 0.5 * (boundbxs_x1 + boundbxs_x2)
|
||
center_y = 0.5 * (boundbxs_y1 + boundbxs_y2)
|
||
return center_x, center_y
|
||
|
||
|
||
def get_rectangle_dimensions(boundbxs):
|
||
'''
|
||
根据检测矩形框，得到其矩形长度和宽度
|
||
输入：两个对角坐标xyxy
|
||
输出：矩形框四个角点坐标，以contours顺序。
|
||
'''
|
||
# 计算两点之间的水平距离
|
||
width = math.fabs(boundbxs[2] - boundbxs[0])
|
||
# 计算两点之间的垂直距离
|
||
height = math.fabs(boundbxs[3]- boundbxs[1])
|
||
return width, height
|
||
|
||
|
||
def fourcorner_coordinate(boundbxs):
|
||
'''
|
||
通过矩形框对角xyxy坐标，得到矩形框轮廓
|
||
输入：两个对角坐标xyxy
|
||
输出：矩形框四个角点坐标，以contours顺序。
|
||
'''
|
||
boundbxs_x1 = boundbxs[0]
|
||
boundbxs_y1 = boundbxs[1]
|
||
boundbxs_x2 = boundbxs[2]
|
||
boundbxs_y2 = boundbxs[3]
|
||
wid = boundbxs_x2 - boundbxs_x1
|
||
hei = boundbxs_y2 - boundbxs_y1
|
||
boundbxs_x3 = boundbxs_x1 + wid
|
||
boundbxs_y3 = boundbxs_y1
|
||
boundbxs_x4 = boundbxs_x1
|
||
boundbxs_y4 = boundbxs_y1 + hei
|
||
contours_rec = [[boundbxs_x1, boundbxs_y1], [boundbxs_x3, boundbxs_y3], [boundbxs_x2, boundbxs_y2],
|
||
[boundbxs_x4, boundbxs_y4]]
|
||
return contours_rec
|
||
|
||
def point_in_rectangle(preds_flag,preds_boat_flag):
|
||
'''
|
||
遍历所有国旗坐标，判断落在检测框中的数量
|
||
输入：
|
||
preds_flag 国旗类别的检测结果列表
|
||
preds_boat_flag 待判定船只的检测结果（单个船只）
|
||
输出：落入检测框的国旗数量
|
||
'''
|
||
iii=0
|
||
boat_contour=fourcorner_coordinate(preds_boat_flag)
|
||
boat_contour=np.array(boat_contour,dtype=np.float32)
|
||
for i in range(len(preds_flag)):
|
||
center_x, center_y = center_coordinate(preds_flag[i])
|
||
if cv2.pointPolygonTest(boat_contour, (center_x, center_y), False)==1:
|
||
iii+=1
|
||
else:
|
||
pass
|
||
return iii
|
||
|
||
|
||
def expand_rect(preds_boat, recScale, imgSize):
|
||
'''
|
||
在图像范围内，将矩形框扩大recScale倍数。
|
||
输入：
|
||
preds_boat 国旗类别的检测结果列表 xyxy_conf_class
|
||
imgSize 从pars传来的元组
|
||
输出：调整后的preds_boat
|
||
'''
|
||
# preds_boat_1=preds_boat
|
||
preds_boat_1=[x for x in preds_boat]
|
||
x1, y1 = preds_boat[0],preds_boat[1]
|
||
x2, y2 = preds_boat[2],preds_boat[3]
|
||
|
||
width = x2 - x1
|
||
height = y2 - y1
|
||
|
||
# 计算新的宽度和高度
|
||
new_width = width * recScale
|
||
new_height = height * recScale
|
||
|
||
# 计算新的对角坐标
|
||
new_x1 = max(x1 - (new_width - width) / 2, 0) # 确保不会超出左边界
|
||
new_y1 = max(y1 - (new_height - height) / 2, 0) # 确保不会超出上边界
|
||
new_x2 = min(x2 + (new_width - width) / 2, imgSize[0]) # 图像宽度是imgSize[0]
|
||
new_y2 = min(y2 + (new_height - height) / 2, imgSize[1]) # 图像高度是imgSize[1]
|
||
|
||
preds_boat_1[0]=new_x1
|
||
preds_boat_1[1]=new_y1
|
||
preds_boat_1[2]=new_x2
|
||
preds_boat_1[3]=new_y2
|
||
|
||
return preds_boat_1
|
||
|
||
|
||
|
||
|
||
###jcq : 增加封仓后处理函数
|
||
def filter_detection_results(results, par):
|
||
target_cls = par['target_cls']
|
||
filter_cls = par['filter_cls']
|
||
|
||
class_target_boxes = [result for result in results if result[5] == target_cls]
|
||
class_filter_boxes = [result for result in results if result[5] == filter_cls]
|
||
filtered_results = []
|
||
|
||
for box_filter in class_filter_boxes:
|
||
is_inside = False
|
||
for box_target in class_target_boxes:
|
||
# 判断filter_cls的框是否完全在target_cls的框内部
|
||
if (box_filter[0] >= box_target[0] and
|
||
box_filter[1] >= box_target[1] and
|
||
box_filter[2] <= box_target[2] and
|
||
box_filter[3] <= box_target[3]):
|
||
is_inside = True
|
||
break
|
||
if is_inside:
|
||
filtered_results.append(box_filter)
|
||
|
||
# 保留所有的target_cls的框以及符合条件的filter_cls的框
|
||
filtered_results += [result for result in class_target_boxes]
|
||
|
||
# 将类别为4的结果映射为5
|
||
for i, result in enumerate(filtered_results):
|
||
if result[5] == 4:
|
||
# 将列表转换为元组，连接后再转换回列表
|
||
new_result = list(tuple(result[:-1]) + (5,))
|
||
filtered_results[i] = new_result
|
||
|
||
return filtered_results
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
|
||
# 对应DJI_20230306140129_0001_Z_165.jpg检测结果
|
||
# preds=[[6.49000e+02, 2.91000e+02, 1.07900e+03, 7.33000e+02, 9.08165e-01, 3.00000e+00],
|
||
# [8.11000e+02, 2.99000e+02, 1.31200e+03, 7.65000e+02, 8.61268e-01, 3.00000e+00],
|
||
# [7.05000e+02, 1.96000e+02, 7.19000e+02, 2.62000e+02, 5.66877e-01, 0.00000e+00]]
|
||
|
||
|
||
# 对应DJI_20230306152702_0001_Z_562.jpg检测结果
|
||
preds=[[7.62000e+02, 7.14000e+02, 1.82800e+03, 9.51000e+02, 9.00902e-01, 3.00000e+00],
|
||
[2.00000e+01, 3.45000e+02, 1.51300e+03, 6.71000e+02, 8.81440e-01, 3.00000e+00],
|
||
[8.35000e+02, 8.16000e+02, 8.53000e+02, 8.30000e+02, 7.07651e-01, 0.00000e+00],
|
||
[1.35600e+03, 4.56000e+02, 1.42800e+03, 4.94000e+02, 6.70549e-01, 2.00000e+00]]
|
||
print('before :\n ',preds)
|
||
#preds=torch.tensor(preds) #返回的预测结果
|
||
imgwidth=1920
|
||
imgheight=1680
|
||
pars={'imgSize':(imgwidth,imgheight),'wRation':1/6.0,'hRation':1/6.0,'smallId':0,'bigId':3,'newId':4,'recScale':1.2}
|
||
# 'smallId':0(国旗)，'bigId':3(船只),wRation和hRation表示判断的阈值条件，newId--新目标的id
|
||
yyy=channel2_post_process([preds],pars) #送入后处理函数
|
||
|
||
print('after :\n ',yyy)
|
||
|
||
|
||
|
||
|
||
|
||
|