| 124 | |
| 125 | |
| 126 | def load_image(image: Union[str, Image.Image], |
| 127 | format: str = "pt", |
| 128 | device: str = "cpu") -> Union[Image.Image, torch.Tensor]: |
| 129 | assert format in ["pt", "pil"], "format must be either Pytorch or PIL" |
| 130 | |
| 131 | if isinstance(image, Image.Image): |
| 132 | return image.convert('RGB') |
| 133 | |
| 134 | parsed_url = urlparse(image) |
| 135 | |
| 136 | if parsed_url.scheme in ["http", "https"]: |
| 137 | image = requests.get(image, stream=True, timeout=10).raw |
| 138 | image = _load_and_convert_image(image) |
| 139 | elif parsed_url.scheme == "data": |
| 140 | image = load_base64_image(parsed_url) |
| 141 | else: |
| 142 | image = _load_and_convert_image(image) |
| 143 | |
| 144 | if format == "pt": |
| 145 | return ToTensor()(image).to(device=device) |
| 146 | else: |
| 147 | return image |
| 148 | |
| 149 | |
| 150 | async def async_load_image( |