test_converter.py 文件源码

def test_image_input(self):
        from _torch_converter import convert
        coreml_model = convert(
            self.model,
            [self.input.shape],
            input_names=['image'],
            image_input_names=['image'],
            preprocessing_args={
                'is_bgr': False,
                'red_bias': 0.0,
                'green_bias': 0.0,
                'blue_bias': 0.0,
                'image_scale': 0.5
            }
        )

        input_array = (np.random.rand(224, 224, 3) * 255).astype('uint8')
        input_image = Image.fromarray(input_array).convert('RGBA')
        output_array = coreml_model.predict({"image": input_image})["output"]
        output_array = output_array.transpose((1, 2, 0))
        npt.assert_array_equal(output_array, input_array * 0.5)