落水人员检测
Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

282 lines
12KB

  1. '''
  2. 这个版本增加了船舶过滤功能
  3. '''
  4. import time
  5. import sys
  6. from core.models.bisenet import BiSeNet
  7. from models.AIDetector_pytorch import Detector
  8. from models.AIDetector_pytorch import plot_one_box,Colors
  9. from utils.postprocess_utils import center_coordinate,fourcorner_coordinate,remove_simivalue,remove_sameeleme_inalist
  10. import os
  11. os.environ['CUDA_VISIBLE_DEVICES'] = '1'
  12. from models.model_stages import BiSeNet
  13. import cv2
  14. import torch
  15. import torch.nn.functional as F
  16. from PIL import Image
  17. import numpy as np
  18. import torchvision.transforms as transforms
  19. from utils.segutils import colour_code_segmentation
  20. from utils.segutils import get_label_info
  21. os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'
  22. os.environ["CUDA_VISIBLE_DEVICES"] = "0"
  23. sys.path.append("../") # 为了导入上级目录的,添加一个新路径
  24. def AI_postprocess(preds,_mask_cv,pars,_img_cv):
  25. '''考虑船上人过滤'''
  26. '''输入:落水人员的结果(类别+坐标)、原图、mask图像
  27. 过程:获得mask的轮廓,判断人员是否在轮廓内。
  28. 在,则保留且绘制;不在,舍弃。
  29. 返回:最终绘制的结果图、最终落水人员(坐标、类别、置信度),
  30. '''
  31. '''1、最大分割水域作为判断依据'''
  32. zoom_factor=4 #缩小因子设置为4,考虑到numpy中分别遍历xy进行缩放耗时大。
  33. original_height = _mask_cv.shape[0]
  34. original_width=_mask_cv.shape[1]
  35. zoom_height=int(original_height/zoom_factor)
  36. zoom_width=int(original_width/zoom_factor)
  37. _mask_cv = cv2.resize(_mask_cv, (zoom_width,zoom_height)) #缩小原图,宽在前,高在后
  38. t4 = time.time()
  39. img_gray = cv2.cvtColor(_mask_cv, cv2.COLOR_BGR2GRAY) if len(_mask_cv.shape)==3 else _mask_cv #
  40. t5 = time.time()
  41. contours, thresh = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
  42. # 寻找轮廓(多边界)
  43. contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, 2)
  44. contour_info = []
  45. for c in contours:
  46. contour_info.append((
  47. c,
  48. cv2.isContourConvex(c),
  49. cv2.contourArea(c),
  50. ))
  51. contour_info = sorted(contour_info, key=lambda c: c[2], reverse=True)
  52. t6 = time.time()
  53. '''新增模块::如果水域为空,则返回原图、无落水人员等。'''
  54. if contour_info==[]:
  55. # final_img=_img_cv
  56. final_head_person_filterwater=[]
  57. timeInfos=0
  58. # return final_img, final_head_person_filterwater
  59. return final_head_person_filterwater,timeInfos
  60. else:
  61. max_contour = contour_info[0]
  62. max_contour1=max_contour[0]
  63. max_contour_X=max_contour1[0][0][:]
  64. max_contour=max_contour[0]*zoom_factor# contours恢复原图尺寸
  65. # max_contour=max_contour[0]*zoom_factor# contours恢复原图尺寸
  66. print(max_contour)
  67. t7 = time.time()
  68. '''2.1、preds中head+person取出,boat取出。'''
  69. init_head_person=[]
  70. init_boat = []
  71. for i in range(len(preds)):
  72. if preds[i][4]=='head' or preds[i][4]=='person':
  73. init_head_person.append(preds[i])
  74. else:
  75. init_boat.append(preds[i])
  76. t8 = time.time()
  77. '''新增模块:2.2、preds中head+person取出,过滤掉head与person中指向同一人的部分,保留同一人的person标签。'''
  78. init_head=[]
  79. init_person=[]
  80. #head与person标签分开
  81. for i in range(len(init_head_person)):
  82. if init_head_person[i][4]=='head':
  83. init_head.append(init_head_person[i])
  84. else:
  85. init_person.append(init_head_person[i])
  86. # person的框形成contours
  87. person_contour=[]
  88. for i in range(len(init_person)):
  89. boundbxs_temp=[init_person[i][0],init_person[i][1],init_person[i][2],init_person[i][3]]
  90. contour_temp_person=fourcorner_coordinate(boundbxs_temp) #得到person预测框的顺序contour
  91. contour_temp_person=np.array(contour_temp_person)
  92. contour_temp_person=np.float32(contour_temp_person)
  93. person_contour.append(np.array(contour_temp_person))
  94. # head是否在person的contours内,在说明是同一人,过滤掉。
  95. list_head=[]
  96. for i in range(len(init_head)):
  97. for j in range(len(person_contour)):
  98. center_x, center_y=center_coordinate(init_head[i])
  99. flag = cv2.pointPolygonTest(person_contour[j], (center_x, center_y), False) #若为False,会找点是否在内,外,或轮廓上(相应返回+1, -1, 0)。
  100. if flag==1:
  101. pass
  102. else:
  103. list_head.append(init_head[i])
  104. # person和最终head合并起来
  105. init_head_person_temp=init_person+list_head
  106. '''3、preds中head+person,通过1中水域过滤'''
  107. init_head_person_filterwater=init_head_person_temp
  108. final_head_person_filterwater=[]
  109. for i in range(len(init_head_person_filterwater)):
  110. center_x, center_y=center_coordinate(init_head_person_filterwater[i])
  111. flag = cv2.pointPolygonTest(max_contour, (center_x, center_y), False) #若为False,会找点是否在内,外,或轮廓上(相应返回+1, -1, 0)。
  112. if flag==1:
  113. final_head_person_filterwater.append(init_head_person_filterwater[i])
  114. else:
  115. pass
  116. t9 = time.time()
  117. '''4、水域过滤后的head+person,再通过船舶范围过滤'''
  118. init_head_person_filterboat=final_head_person_filterwater
  119. # final_head_person_filterboat=[]
  120. #获取船舶范围
  121. boat_contour=[]
  122. for i in range(len(init_boat)):
  123. boundbxs1=[init_boat[i][0],init_boat[i][1],init_boat[i][2],init_boat[i][3]]
  124. contour_temp=fourcorner_coordinate(boundbxs1) #得到boat预测框的顺序contour
  125. contour_temp_=np.array(contour_temp)
  126. contour_temp_=np.float32(contour_temp_)
  127. boat_contour.append(np.array(contour_temp_))
  128. t10 = time.time()
  129. # 遍历船舶范围,取出在船舶范围内的head和person(可能有重复元素)
  130. list_headperson_inboat=[]
  131. for i in range(len(init_head_person_filterboat)):
  132. for j in range(len(boat_contour)):
  133. center_x, center_y=center_coordinate(init_head_person_filterboat[i])
  134. # yyyyyyyy=boat_contour[j]
  135. flag = cv2.pointPolygonTest(boat_contour[j], (center_x, center_y), False) #若为False,会找点是否在内,外,或轮廓上(相应返回+1, -1, 0)。
  136. if flag==1:
  137. list_headperson_inboat.append(init_head_person_filterboat[i])
  138. else:
  139. pass
  140. print('list_headperson_inboat',list_headperson_inboat)
  141. if len(list_headperson_inboat)==0:
  142. pass
  143. else:
  144. list_headperson_inboat=remove_sameeleme_inalist(list_headperson_inboat) #将重复嵌套列表元素删除
  145. # 过滤船舶范围内的head和person
  146. final_head_person_filterboat=remove_simivalue(init_head_person_filterboat,list_headperson_inboat)
  147. t11 = time.time()
  148. '''5、输出最终落水人员,并绘制保存检测图'''
  149. colors = Colors()
  150. if final_head_person_filterwater is not None:
  151. for i in range(len(final_head_person_filterboat)):
  152. # lbl = self.names[int(cls_id)]
  153. lbl = final_head_person_filterboat[i][4]
  154. xyxy=[final_head_person_filterboat[i][0],final_head_person_filterboat[i][1],final_head_person_filterboat[i][2],final_head_person_filterboat[i][3]]
  155. c = int(5)
  156. plot_one_box(xyxy, _img_cv, label=lbl, color=colors(c, True), line_thickness=3)
  157. final_img=_img_cv
  158. t12 = time.time()
  159. # cv2.imwrite('final_result.png', _img_cv)
  160. t13 = time.time()
  161. print('存图:%s, 过滤标签:%s ,遍历船舶范围:%s,水域过滤后的head+person:%s,水域过滤:%s,head+person、boat取出:%s,新增如果水域为空:%s,找contours:%s,图像改变:%s'
  162. %((t13-t12) * 1000,(t12-t11) * 1000,(t11-t10) * 1000,(t10-t9) * 1000,(t9-t8) * 1000,(t8-t7) * 1000,(t7-t6) * 1000,(t6-t5) * 1000,(t5-t4) * 1000 ) )
  163. timeInfos=('存图:%s, 过滤标签:%s ,遍历船舶范围:%s,水域过滤后的head+person:%s,水域过滤:%s,head+person、boat取出:%s,新增如果水域为空:%s,找contours:%s,图像改变:%s'
  164. %((t13-t12) * 1000,(t12-t11) * 1000,(t11-t10) * 1000,(t10-t9) * 1000,(t9-t8) * 1000,(t8-t7) * 1000,(t7-t6) * 1000,(t6-t5) * 1000,(t5-t4) * 1000 ) )
  165. return final_head_person_filterwater,timeInfos #返回最终绘制的结果图、最终落水人员(坐标、类别、置信度)
  166. def AI_process(model, segmodel, args1,path1):
  167. '''对原图进行目标检测和水域分割'''
  168. '''输入:检测模型、分割模型、配置参数、路径
  169. 返回:返回目标检测结果、原图像、分割图像,
  170. '''
  171. '''检测图片'''
  172. t21=time.time()
  173. _img_cv = cv2.imread(path1) # 将这里的送入yolov5
  174. t22 = time.time()
  175. # _img_cv=_img_cv.numpy()
  176. pred = model.detect(_img_cv) # 检测结果
  177. #对pred处理,处理成list嵌套
  178. pred=[[*x[0:4],x[4],x[5].cpu().tolist()] for x in pred[1]]
  179. # pred=[[x[0],*x[1:5],x[5].cpu().float()] for x in pred[1]]
  180. print('pred', pred)
  181. t23 = time.time()
  182. '''分割图片'''
  183. img = Image.open(path1).convert('RGB')
  184. t231 = time.time()
  185. transf1 = transforms.ToTensor()
  186. transf2 = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
  187. imgs = transf1(img)
  188. imgs = transf2(imgs)
  189. print(path1) # numpy数组格式为(H,W,C)
  190. size = [360, 640]
  191. imgs = imgs.unsqueeze(0)
  192. imgs = imgs.cuda()
  193. N, C, H, W = imgs.size()
  194. self_scale = 360 / H
  195. new_hw = [int(H * self_scale), int(W * self_scale)]
  196. print("line50", new_hw)
  197. imgs = F.interpolate(imgs, new_hw, mode='bilinear', align_corners=True)
  198. t24 = time.time()
  199. with torch.no_grad():
  200. logits = segmodel(imgs)[0]
  201. t241 = time.time()
  202. logits = F.interpolate(logits, size=size, mode='bilinear', align_corners=True)
  203. probs = torch.softmax(logits, dim=1)
  204. preds = torch.argmax(probs, dim=1)
  205. preds_squeeze = preds.squeeze(0)
  206. preds_squeeze_predict = colour_code_segmentation(np.array(preds_squeeze.cpu()), args1['label_info'])
  207. preds_squeeze_predict = cv2.resize(np.uint8(preds_squeeze_predict), (W, H))
  208. predict_mask = cv2.cvtColor(np.uint8(preds_squeeze_predict), cv2.COLOR_RGB2BGR)
  209. _mask_cv =predict_mask
  210. t25 = time.time()
  211. cv2.imwrite('seg_result.png', _mask_cv)
  212. t26 = time.time()
  213. print('存分割图:%s, 分割后处理:%s ,分割推理:%s ,分割图变小:%s,分割图读图:%s,检测模型推理:%s,读图片:%s'
  214. %((t26-t25) * 1000,(t25-t241) * 1000,(t241-t24) * 1000,(t24-t231) * 1000,(t231-t23) * 1000,(t23-t22) * 1000,(t22-t21) * 1000 ) )
  215. return pred, _img_cv, _mask_cv #返回目标检测结果、原图像、分割图像
  216. def main():
  217. '''配置参数'''
  218. label_info = get_label_info('utils/class_dict.csv')
  219. pars={'cuda':'0','crop_size':512,'input_dir':'input_dir','output_dir':'output_dir','workers':16,'label_info':label_info,
  220. 'dspth':'./data/','backbone':'STDCNet813','use_boundary_2':False, 'use_boundary_4':False, 'use_boundary_8':True, 'use_boundary_16':False,'use_conv_last':False}
  221. dete_weights='weights/best_luoshui20230608.pt'
  222. '''分割模型权重路径'''
  223. seg_weights = 'weights/model_final.pth'
  224. '''初始化目标检测模型'''
  225. model = Detector(dete_weights)
  226. '''初始化分割模型2'''
  227. n_classes = 2
  228. segmodel = BiSeNet(backbone=pars['backbone'], n_classes=n_classes,
  229. use_boundary_2=pars['use_boundary_2'], use_boundary_4=pars['use_boundary_4'],
  230. use_boundary_8=pars['use_boundary_8'], use_boundary_16=pars['use_boundary_16'],
  231. use_conv_last=pars['use_conv_last'])
  232. segmodel.load_state_dict(torch.load(seg_weights))
  233. segmodel.cuda()
  234. segmodel.eval()
  235. '''图像测试'''
  236. folders = os.listdir(pars['input_dir'])
  237. for i in range(len(folders)):
  238. path1 = pars['input_dir'] + '/' + folders[i]
  239. t1=time.time()
  240. '''对原图进行目标检测和水域分割'''
  241. pred, _img_cv, _mask_cv=AI_process(model,segmodel, pars,path1)
  242. t2 = time.time()
  243. '''进入后处理,判断水域内有落水人员'''
  244. haha,zzzz=AI_postprocess(pred, _mask_cv,pars,_img_cv )
  245. t3 = time.time()
  246. print('总时间分布:前处理t2-t1,后处理t3-t2',(t2-t1)*1000,(t3-t2)*1000)
  247. if __name__ == "__main__":
  248. main()