File size: 2,248 Bytes
7cdc7d0 d08fbc6 7cdc7d0 d08fbc6 7cdc7d0 d08fbc6 7cdc7d0 d08fbc6 7cdc7d0 d08fbc6 7cdc7d0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import base64
import io
import re
from abc import abstractmethod
from typing import Any, Dict
import numpy as np
from .dict_utils import dict_get
from .operators import FieldOperator, InstanceFieldOperator, PackageRequirementsMixin
class PillowMixin(PackageRequirementsMixin):
_requirements_list = {"PIL": "pip install pillow"}
def prepare(self):
super().prepare()
import PIL
from PIL import Image
self.pil = PIL
self.image = Image
def extract_images(text, instance):
regex = r'<img\s+src=["\'](.*?)["\']'
image_sources = re.findall(regex, text)
images = []
for image_source in image_sources:
image = dict_get(instance, image_source)
images.append(image)
return images
class DecodeImage(FieldOperator, PillowMixin):
def decode_base64_to_image(self, base64_string):
image_data = base64.b64decode(base64_string)
return self.image.open(io.BytesIO(image_data))
def process_value(self, value: Any) -> Any:
return {"image": self.decode_base64_to_image(value)}
class ToImage(InstanceFieldOperator):
def process_instance_value(self, value: Any, instance: Dict[str, Any]):
return {"image": value}
class ImageFieldOperator(FieldOperator, PillowMixin):
@abstractmethod
def process_image(self, image):
pass
def process_value(self, value: Any) -> Any:
if not isinstance(value, self.image.Image):
raise ValueError(f"ImageFieldOperator requires image, got {type(value)}.")
return self.process_image(value)
class GrayScale(ImageFieldOperator):
def process_image(self, image):
# Convert the image to grayscale
grayscale_image = image.convert("L")
# Convert the grayscale image to a NumPy array
grayscale_array = np.array(grayscale_image)
# Add a dummy channel dimension to make it (height, width, 1)
grayscale_array = np.expand_dims(grayscale_array, axis=-1)
# Repeat the channel to have (height, width, 3) if needed for compatibility
grayscale_array = np.repeat(grayscale_array, 3, axis=-1)
# Convert back to a PIL image with 3 channels
return self.image.fromarray(grayscale_array)
|