| 17 | |
| 18 | @register(inputs=['inp'], outputs=['out']) |
| 19 | class Detect: |
| 20 | def __init__(self, name, args): |
| 21 | logger.info("loading YOLOX detection...") |
| 22 | self._tsize = args['tsize'] |
| 23 | self._interval = args['interval'] |
| 24 | self._visualize = args['visualize'] |
| 25 | self.name = name |
| 26 | |
| 27 | # load detect model and warmup |
| 28 | self._predictor = PredictorLite(path=args['path'], |
| 29 | confthre=args['conf'], |
| 30 | nmsthre=args['nms'], |
| 31 | test_size=(self._tsize, self._tsize), |
| 32 | device=args['device'], |
| 33 | device_id=args['device_id']) |
| 34 | warmup_data = np.zeros((224, 224, 3), dtype=np.uint8) |
| 35 | self._predictor.inference(warmup_data) |
| 36 | logger.info(" YOLOX loaded.") |
| 37 | |
| 38 | @staticmethod |
| 39 | def restrict(val, min, max): |
| 40 | assert min < max |
| 41 | if val < min: |
| 42 | val = min |
| 43 | if val > max: |
| 44 | val = max |
| 45 | return val |
| 46 | |
| 47 | def exec(self): |
| 48 | envelope = self.inp.recv() |
| 49 | if envelope is None: |
| 50 | return |
| 51 | |
| 52 | msg = envelope.msg |
| 53 | msg['items'] = [] |
| 54 | |
| 55 | process = envelope.partial_id % self._interval == 0 |
| 56 | if process: |
| 57 | data = msg['data'] |
| 58 | outputs = self._predictor.inference(data) |
| 59 | # skip if detect nothing |
| 60 | if outputs is not None: |
| 61 | items = [] |
| 62 | |
| 63 | for i in range(outputs.shape[0]): |
| 64 | output = outputs[i] |
| 65 | item = dict() |
| 66 | item["bbox"] = output[0:4] |
| 67 | item["det_score"] = output[4] * output[5] |
| 68 | items.append(item) |
| 69 | msg['items'] = items |
| 70 | |
| 71 | # import cv2 |
| 72 | # x = self._predictor.visual(outputs, data) |
| 73 | # name = 'frame{0:07d}.jpg'.format(envelope.partial_id) |
| 74 | # cv2.imwrite(name, x) |
| 75 | |
| 76 | if self._visualize == 1: |
nothing calls this directly
no outgoing calls
no test coverage detected