diff --git a/README.md b/README.md index fc11237..2c08185 100644 --- a/README.md +++ b/README.md @@ -18,11 +18,11 @@ In computer vision (CV) directions, image transformations are usually an indispe ​ The size of the output image. -`__call__(self, img_tensor: Union[np.ndarray, Image.Image, torch.Tensor, str])` +`__call__(self, img: Union[np.ndarray, Image.Image, torch.Tensor, str])` **Args:** -​ img_tensor(`Union[np.ndarray, Image.Image, torch.Tensor, str]`): +​ img(`Union[np.ndarray, Image.Image, torch.Tensor, str]`): ​ Original image data, the type can be np.ndarry, PIL.image, or str path of the image. @@ -38,7 +38,7 @@ You can get the required python package by [requirements.txt](./requirements.txt ## How it works -The `towhee/transform-image` Operator is used for image transformation and is an important part of data preprocessing. It can be added to the pipeline and is usually used as the first custom operator of the pipeline. For example, it's the first Operator named processing within [image_embedding_resnet50](https://hub.towhee.io/towhee/image-embedding-resnet50) pipeline, and it is the red box in the picture below. +The `towhee/transform-image` Operator is used for image transformation and is an important part of data preprocessing. It can be added to the pipeline and is usually used as the first custom operator of the pipeline. For example, it's the first Operator named processing within [image-embedding-resnet50](https://hub.towhee.io/towhee/image-embedding-resnet50) pipeline, and it is the red box in the picture below. ![img](./readme_src/operator.png) diff --git a/transform_image.py b/transform_image.py index 656832b..04703c2 100644 --- a/transform_image.py +++ b/transform_image.py @@ -43,21 +43,21 @@ class TransformImage(Operator): ] ) - def __call__(self, img_tensor: Union[np.ndarray, Image.Image, torch.Tensor, str]) -> NamedTuple('Outputs', [('img_transformed', torch.Tensor)]): + def __call__(self, img: Union[np.ndarray, Image.Image, torch.Tensor, str]) -> NamedTuple('Outputs', [('img_transformed', torch.Tensor)]): """ Call it when use this class. Args: - img_tensor (`Union[np.ndarray, Image.Image, torch.Tensor, str]`): + img(`Union[np.ndarray, Image.Image, torch.Tensor, str]`): The image data to be normalized, you can try one of the four formats: np.ndarray, Image.Image, torch.Tensor and str. Returns: (`torch.Tensor`) The normalized image tensor. """ - if isinstance(img_tensor, str): - img_tensor = Image.open(img_tensor) - if isinstance(img_tensor, Image.Image): - img_tensor = img_tensor.convert('RGB') + if isinstance(img, str): + img_tensor = Image.open(img) + if isinstance(img, Image.Image): + img_tensor = img.convert('RGB') Outputs = NamedTuple('Outputs', [('img_transformed', torch.Tensor)]) return Outputs(self.tfms(img_tensor).unsqueeze(0)) diff --git a/transform_image.yaml b/transform_image.yaml index 1d8734f..3d1cd2d 100644 --- a/transform_image.yaml +++ b/transform_image.yaml @@ -8,6 +8,6 @@ init: size: int call: input: - img_tensor: Union[np.ndarray, Image.Image, torch.Tensor, str] + img: Union[np.ndarray, Image.Image, torch.Tensor, str] output: img_transformed: torch.Tensor \ No newline at end of file