YOLOX postprocessing

[umairjavaid]

How can I implement this post processing step in FastMOT

def yolox_nano_postprocess(self, prediction, num_classes=2, conf_thre=0.01, nms_thre=0.65, class_agnostic=True):
        print("#" * 100)
        box_corner = prediction.copy()
        print("box corner shape : {}".format(box_corner.shape))
        box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
        box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
        box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
        box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
        prediction[:, :, :4] = box_corner[:, :, :4]
        print("prediction shape : {}".format(prediction.shape))
​
        output = [None for _ in range(len(prediction))]
        print("output shape : {}".format(np.asarray(output).shape))
        for i, image_pred in enumerate(prediction):
            print("VALUE OF I : {}".format(i))
            # If none are remaining => process next image
            if not image_pred.shape[0]:
                continue
​
            print("image-pred shape : {}".format(image_pred.shape))
            # Get score and class with highest confidence
            print("image_pred[:, 5: 5 + num_classes] shape : {}".format(image_pred[:, 5: 5 + num_classes].shape))
            print("image_pred[:, 5: 5 + num_classes] type : {}".format(type(image_pred[:, 5: 5 + num_classes])))
            # class_conf, class_pred = np.max(image_pred[:, 5: 5 + num_classes], axis=0)
            class_conf_pred = image_pred[:, 5: 5 + num_classes]
            class_conf, class_pred = class_conf_pred[:, 0], class_conf_pred[:, 1]
            class_conf = class_conf.reshape(class_conf.shape[0], 1)
            class_pred = class_pred.reshape(class_pred.shape[0], 1)
            print("class conf shape : {}".format(class_conf.shape))
            print("calss pred shape : {}".format(class_pred.shape))
​
            conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_thre).squeeze()
            # conf_mask = conf_mask.reshape(conf_mask.shape[0], 1)
            print("conf mask shape : {}".format(conf_mask.shape))
            # Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred)
            print("image_pred[:, :5] shape : {}".format(image_pred[:, :5].shape))
            detections = np.concatenate((image_pred[:, :5], class_conf, class_pred.astype(float)), axis=1)
            print("detections shape : {}".format(detections.shape))
            detections = detections[conf_mask]
            print("detections2 shape : {}".format(detections.shape))
            if not detections.shape[0]:
                continue
​
            if True: #class_agnostic:
                nms_out_index = self.nms(detections[:, :4], detections[:, 4] * detections[:, 5], nms_thre)
                # print("nms out index in class agnostic true : {}".format(nms_out_index))
                # print("nms out index 0 in class agnostic true shape : {}".format(nms_out_index[0].shape))
                # print("nms out index 0 in class agnostic true shape : {}".format(nms_out_index[1].shape))
                # print("nms out index 0 in class agnostic true shape : {}".format(nms_out_index[2].shape))
​
                # detections = detections[nms_out_index]
                # print("detections afteer nms shape : {}".format(detections.shape))
​
                bboxes , confidences, classes = nms_out_index
                dets = []
                for bbox, conf, pred_cls in zip(bboxes , confidences, classes):
                    # print(type(bbox), type(conf[0]), type(pred_cls))
                    detection = [bbox, conf[0], int(pred_cls)]
                    # print(type(detection), detection)
                    dets.append(detection)
            if output[i] is None:
                output[i] = np.asarray(dets)
            else:
                output[i] = np.concatenate((output[i], np.asarray(dets)))
​
            print("output after nms shape : {}".format(np.asarray(output).shape))
​
        print("#" * 100)
​
        return output