def do_transform_ip()

in awstreamer/gst_plugins/python/neodlr.py [0:0]

• 38 lines of code
• 8 McCabe index (conditional complexity)

    def do_transform_ip(self, buffer: Gst.Buffer) -> Gst.FlowReturn:
        try:
            # Load the model
            if self.model is None:
                self.load_model()

            # Check if model has been loaded
            if self.model is None:
                return Gst.FlowReturn.OK

            # convert Gst.Buffer to np.ndarray
            image = gst_buffer_with_caps_to_ndarray(buffer, self.sinkpad.get_current_caps())

            print('Testing inference...')
            start_time = time.time()
            # img_rand = np.random.rand(1, 3, 320, 320)

            # Prepare input
            image_3 = image[:,:,:3]
            img_small = cv2.resize(image_3, (self.image_size, self.image_size))
            img_reshaped = np.reshape(img_small, (1, 3, self.image_size, self.image_size))

            # # Normalize & transpose
            # mean_vec = np.array([0.485, 0.456, 0.406])
            # stddev_vec = np.array([0.229, 0.224, 0.225])
            # img_reshaped = (img_reshaped/255 - mean_vec)/stddev_vec
            # img_reshaped = np.rollaxis(img_reshaped, axis=2, start=0)[np.newaxis, :]

            # Run inference
            result = self.model.run(img_reshaped)
            print('inference time is ' + str((time.time()-start_time)) + ' seconds')

            # Process inference output
            temp = []
            for r in result:
                r = np.squeeze(r)
                temp.append(r.tolist())
            idx, score, bbox = temp
            bbox = np.asarray(bbox)
            res = np.hstack((np.column_stack((idx, score)), bbox))
            l = list()
            for r in res:
                (class_id, score, x0, y0, x1, y1) = r
                if score < self.threshold:
                    continue
                d = {
                    "bounding_box": (int(x0), int(y0), int(x1), int(y1)),
                    "confidence": score,
                    "class_name": "class_name",
                    "track_id": int(class_id)
                }
                l.append(d)
            print(l)
            gst_meta_write(buffer, l)

        except Exception as e:
            logging.error(e)

        return Gst.FlowReturn.OK