train.py 文件源码

 复制代码def __call__(self, in_data):
        orig_img, bboxes, whole_mask, labels = in_data
        _, orig_H, orig_W = orig_img.shape
        img = self.model.prepare(
            orig_img, self.target_height, self.max_width)
        img = img.astype(np.float32)
        del orig_img
        _, H, W = img.shape
        scale = H / orig_H

        bboxes = chainercv.transforms.resize_bbox(
            bboxes, (orig_H, orig_W), (H, W))

        indices = get_keep_indices(bboxes)
        bboxes = bboxes[indices, :]
        whole_mask = whole_mask[indices, :, :]
        labels = labels[indices]

        whole_mask = whole_mask.transpose((1, 2, 0))
        whole_mask = cv2.resize(
            whole_mask.astype(np.uint8), (W, H),
            interpolation=cv2.INTER_NEAREST)
        if whole_mask.ndim < 3:
            whole_mask = whole_mask.reshape((H, W, 1))
        whole_mask = whole_mask.transpose((2, 0, 1))

        if self.flip:
            img, params = chainercv.transforms.random_flip(
                img, x_random=True, return_param=True)
            whole_mask = fcis.utils.flip_mask(
                whole_mask, x_flip=params['x_flip'])
            bboxes = chainercv.transforms.flip_bbox(
                bboxes, (H, W), x_flip=params['x_flip'])

        return img, bboxes, whole_mask, labels, scale