r""" Construct a GeoLayoutLM image processor Args: do_preprocess (`bool`): whether to do preprocess to unify the image format, resize and convert to tensor. do_rescale: only works when we disable do_preprocess.
| 23 | |
| 24 | |
| 25 | class ImageProcessor(object): |
| 26 | r""" |
| 27 | Construct a GeoLayoutLM image processor |
| 28 | Args: |
| 29 | do_preprocess (`bool`): whether to do preprocess to unify the image format, |
| 30 | resize and convert to tensor. |
| 31 | do_rescale: only works when we disable do_preprocess. |
| 32 | """ |
| 33 | |
| 34 | def __init__(self, |
| 35 | do_preprocess: bool = True, |
| 36 | do_resize: bool = False, |
| 37 | image_size: Dict[str, int] = None, |
| 38 | do_rescale: bool = False, |
| 39 | rescale_factor: float = 1. / 255, |
| 40 | do_normalize: bool = True, |
| 41 | image_mean: Union[float, Iterable[float]] = None, |
| 42 | image_std: Union[float, Iterable[float]] = None, |
| 43 | apply_ocr: bool = True, |
| 44 | **kwargs) -> None: |
| 45 | self.do_preprocess = do_preprocess |
| 46 | self.do_resize = do_resize |
| 47 | self.size = image_size if image_size is not None else { |
| 48 | 'height': 768, |
| 49 | 'width': 768 |
| 50 | } |
| 51 | self.do_rescale = do_rescale and (not do_preprocess) |
| 52 | self.rescale_factor = rescale_factor |
| 53 | self.do_normalize = do_normalize |
| 54 | image_mean = IMAGENET_DEFAULT_MEAN if image_mean is None else image_mean |
| 55 | image_std = IMAGENET_DEFAULT_STD if image_std is None else image_std |
| 56 | self.image_mean = (image_mean, image_mean, image_mean) if isinstance( |
| 57 | image_mean, float) else image_mean |
| 58 | self.image_std = (image_std, image_std, image_std) if isinstance( |
| 59 | image_std, float) else image_std |
| 60 | self.apply_ocr = apply_ocr |
| 61 | self.kwargs = kwargs |
| 62 | |
| 63 | self.totensor = transforms.ToTensor() |
| 64 | |
| 65 | def preprocess(self, image: Union[np.ndarray, PIL.Image.Image]): |
| 66 | """ unify the image format, resize and convert to tensor. |
| 67 | """ |
| 68 | image = LoadImage.convert_to_ndarray(image)[:, :, ::-1] |
| 69 | size_raw = image.shape[:2] |
| 70 | if self.do_resize: |
| 71 | image = cv2.resize(image, |
| 72 | (self.size['width'], self.size['height'])) |
| 73 | # convert to pytorch tensor |
| 74 | image_pt = self.totensor(image) |
| 75 | return image_pt, size_raw |
| 76 | |
| 77 | def __call__(self, images: Union[list, np.ndarray, PIL.Image.Image, str]): |
| 78 | """ |
| 79 | Args: |
| 80 | images: list of np.ndarrays, PIL images or image tensors. |
| 81 | """ |
| 82 | if not isinstance(images, list): |
no outgoing calls
no test coverage detected
searching dependent graphs…