code: #!/usr/bin/env python
  # -*- coding: UTF-8 -*-
  
  # n(net) o(oil) h(hang) r(rust) 检测模块
  import os
  import sys
  current_dir = os.path.dirname(os.path.abspath(__file__))
  sys.path.append(current_dir)
  
  from mmdet.models import build_detector
  import mmcv
  import torch
  import cv2
  import time
  import json
  from mmcv.runner import load_checkpoint
  import PIL.Image as Image
  import numpy as np
  from torchvision.transforms import transforms
  import pycocotools.mask as maskUtils
  
  current_dir = os.path.dirname(os.path.abspath(__file__))
  config_file = os.path.join(current_dir, 'configs','config_cascade_rcnn.py')
  weight_file = '/home/kilox/weights/nohr_best.pth'
  # weight_file = '/Weights/verified/oil_detection_v1/oil_best.pth'
  
  class Object(object):
      def __init__(self):
          self.class_name = "Unknown"
          self.trust = 0.0
          self.rank = 0
      
      def to_json(self):
          return json.dumps(self.__dict__)
  
  
  class Port:
      def __init__(self):
          self.cfg = mmcv.Config.fromfile(config_file)
          # 创建模型 , test_cfg 是rpn rcnn的nms等配置
          self.detector = build_detector(self.cfg.model, train_cfg=None, test_cfg=self.cfg.test_cfg)
          # 加载权重
          load_checkpoint(self.detector, weight_file, map_location='cpu')
          self.detector = self.detector.to('cuda')
          self.detector.eval()
          self.class_names = ('油污','鸟巢','锈蚀','飘挂物')
      
      def process(self, image,save=None):
          """
          :param image: PIL.Image 输入图像
          """
          np_image = np.asarray(image)
          img, img_meta = self.prepare_single(np_image)
          # forward
          with torch.no_grad():
              # 传入rescale则代表返回的mask是原图的
              result = self.detector.simple_test(img, [img_meta], proposals=None, rescale=True)
              # 将mask 以及bbox画在图上
              img = self.draw_image(np_image, img_meta, result)
              real_time = time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime(time.time()))
              output_file_name = os.path.join(real_time + '.jpg')
              cv2.imwrite(output_file_name, img)
          return False,None,output_file_name
  
      # 将图片添加meta的函数
      def prepare_single(self,img):
          img_info = {'height': img.shape[0], 'width': img.shape[1]}
          img_norm_cfg = self.cfg.img_norm_cfg
          size_divisor = self.cfg.data.test.size_divisor
          
          img, scale_factor = mmcv.imrescale(img, (4014,2400), return_scale=True)
          img_shape = img.shape
          
          img = mmcv.imnormalize(img, img_norm_cfg.mean, img_norm_cfg.std, img_norm_cfg.to_rgb)
          img = mmcv.impad_to_multiple(img, size_divisor)
          pad_shape = img.shape
          _img = transforms.ToTensor()(img).float()
          _img = _img.unsqueeze(0)
          _img_meta = dict(
              ori_shape=(img_info['height'], img_info['width'], 3),
              img_shape=img_shape,
              pad_shape=pad_shape,
              scale_factor=scale_factor,
              flip=False)
          _img = _img.to('cuda')
          return _img, _img_meta,
  
      def draw_image(self,img, meta, result, score_thr=0.9):
          def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True):
              num_imgs = tensor.size(0)
              mean = np.array(mean, dtype=np.float32)
              std = np.array(std, dtype=np.float32)
              imgs = []
              for img_id in range(num_imgs):
                  img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
                  img = mmcv.imdenormalize(
                      img, mean, std, to_bgr=to_rgb).astype(np.uint8)
                  imgs.append(np.ascontiguousarray(img))
              return imgs
          
          if isinstance(result, tuple):
              bbox_result, segm_result = result
          else:
              bbox_result, segm_result = result, None
          
          h, w, _ = meta['ori_shape']
          img_show = img[:h, :w, :].copy()
          
          bboxes = np.vstack(bbox_result)
          # 画mask
          # # draw segmentation masks
          # if segm_result is not None:
          #     segms = mmcv.concat_list(segm_result)
          #     inds = np.where(bboxes[:, -1] > score_thr)[0]
          #     for i in inds:
          #         color_mask = np.random.randint(
          #             0, 256, (1, 3), dtype=np.uint8)
          #         mask = maskUtils.decode(segms[i]).astype(np.bool)
          #         # todo fix dimension not equal
          #         img_check_shape = tuple(img_show.shape[0:2])
          #         if mask.shape != img_check_shape:
          #             width_diff = mask.shape[1] - img_check_shape[1]
          #             if mask.shape[1] < img_check_shape[1]:
          #                 mask = np.pad(mask, (0, width_diff), mode='constant', constant_values=False)
          #                 np.insert(mask, False, )
          #             else:
          #                 mask = mask[:, :-width_diff]
          #         img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
          # 画bbox
          labels = [
              np.full(bbox.shape[0], i, dtype=np.int32)
              for i, bbox in enumerate(bbox_result)
          ]
          labels = np.concatenate(labels)
          assert bboxes.shape[1] == 5
          scores = bboxes[:, -1]
          inds = scores > score_thr
          bboxes = bboxes[inds, :]
          labels = labels[inds]
      
          for bbox, label in zip(bboxes, labels):
              bbox_int = bbox.astype(np.int32)
              left_top = (bbox_int[0], bbox_int[1])
              right_bottom = (bbox_int[2], bbox_int[3])
              cv2.rectangle(
                  img_show, left_top, right_bottom, (0, 255, 0), thickness=2)
              label_text = self.class_names[
                  label] if self.class_names is not None else 'cls {}'.format(label)
              if len(bbox) > 4:
                  label_text += '|{:.02f}'.format(bbox[-1])
              cv2.putText(img_show, label_text, (bbox_int[0], bbox_int[1] - 2),
                          cv2.FONT_HERSHEY_COMPLEX, 2, (0, 255, 0))
          
          return img_show
  
  def test():
      pass
  
  
  if __name__ == '__main__':
      im = Image.open('/home/kilox/3.jpg')
      port = Port()
      print(port.process(im,True))
  
apis: numpy.vstack numpy.ascontiguousarray torch.no_grad numpy.asarray numpy.array numpy.concatenate numpy.full