用kafka接收消息
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

129 line
4.4KB

  1. import torch
  2. from torchvision import transforms
  3. import cv2,os,sys
  4. sys.path.extend(['segutils'])
  5. import numpy as np
  6. import matplotlib.pyplot as plt
  7. from model_stages import BiSeNet
  8. import torch.nn.functional as F
  9. import time
  10. class SegModel(object):
  11. def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0'):
  12. self.model = BiSeNet(backbone='STDCNet813', n_classes=nclass,
  13. use_boundary_2=False, use_boundary_4=False,
  14. use_boundary_8=True, use_boundary_16=False,
  15. use_conv_last=False)
  16. self.device = device
  17. self.model.load_state_dict(torch.load(weights, map_location=torch.device(self.device) ))
  18. self.model= self.model.to(self.device)
  19. self.mean = (0.485, 0.456, 0.406)
  20. self.std = (0.229, 0.224, 0.225)
  21. def eval(self, image=None):
  22. time0 = time.time()
  23. imageH, imageW, _ = image.shape
  24. image = self.RB_convert(image)
  25. img = self.preprocess_image(image)
  26. if self.device != 'cpu':
  27. imgs = img.to(self.device)
  28. else:imgs=img
  29. time1 = time.time()
  30. self.model.eval()
  31. with torch.no_grad():
  32. output = self.model(imgs)
  33. time2 = time.time()
  34. pred = output.data.cpu().numpy()
  35. pred = np.argmax(pred, axis=1)[0]#得到每行
  36. time3 = time.time()
  37. pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
  38. time4 = time.time()
  39. outstr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
  40. return pred, outstr
  41. def preprocess_image(self, image):
  42. image = cv2.resize(image, (640,360), interpolation=cv2.INTER_LINEAR)
  43. image = image.astype(np.float32)
  44. image /= 255.0
  45. image[:, :, 0] -= self.mean[0]
  46. image[:, :, 1] -= self.mean[1]
  47. image[:, :, 2] -= self.mean[2]
  48. image[:, :, 0] /= self.std[0]
  49. image[:, :, 1] /= self.std[1]
  50. image[:, :, 2] /= self.std[2]
  51. image = np.transpose(image, (2, 0, 1))
  52. image = torch.from_numpy(image).float()
  53. image = image.unsqueeze(0)
  54. return image
  55. def get_ms(self,t1,t0):
  56. return (t1-t0)*1000.0
  57. def RB_convert(self,image):
  58. image_c = image.copy()
  59. image_c[:,:,0] = image[:,:,2]
  60. image_c[:,:,2] = image[:,:,0]
  61. return image_c
  62. def get_ms(t1,t0):
  63. return (t1-t0)*1000.0
  64. def get_largest_contours(contours):
  65. areas = [cv2.contourArea(x) for x in contours]
  66. max_area = max(areas)
  67. max_id = areas.index(max_area)
  68. return max_id
  69. if __name__=='__main__':
  70. impth = '/home/thsw/WJ/src/STDC/images/examples'
  71. outpth= '/home/thsw/WJ/src/STDC/images/results'
  72. folders = os.listdir(impth)
  73. weights = '/home/thsw/WJ/src/STDC/model_maxmIOU75_1720_0.946_360640.pth'
  74. segmodel = SegModel(nclass=2,weights=weights)
  75. for i in range(len(folders)):
  76. imgpath = os.path.join(impth, folders[i])
  77. time0 = time.time()
  78. #img = Image.open(imgpath).convert('RGB')
  79. img = cv2.imread(imgpath)
  80. img = np.array(img)
  81. time1 = time.time()
  82. pred, outstr = segmodel.eval(image=img)#####
  83. time2 = time.time()
  84. binary0 = pred.copy()
  85. contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
  86. time3 = time.time()
  87. max_id = -1
  88. if len(contours)>0:
  89. max_id = get_largest_contours(contours)
  90. binary0[:,:] = 0
  91. cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1)
  92. cv2.drawContours(img,contours,max_id,(0,255,255),3)
  93. time4 = time.time()
  94. #img_n = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
  95. cv2.imwrite( os.path.join( outpth,folders[i] ) ,img )
  96. time5 = time.time()
  97. print('image:%d ,infer:%.1f ms,findcontours:%.1f ms, draw:%.1f, total:%.1f'%(i,get_ms(time2,time1),get_ms(time3,time2),get_ms(time4,time3),get_ms(time4,time1)))