Spaces:
Runtime error
Runtime error
| import logging | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ExifTags, ImageCms | |
| from io import BytesIO | |
| import requests | |
| import pillow_heif | |
| from src.segmentation_utils import ( | |
| detect_faces_and_landmarks, | |
| detect_face_landmarks, | |
| mediapipe_selfie_segmentor, | |
| create_feature_masks, | |
| ) | |
| from src.utils import is_url, extract_filename_and_extension | |
| # Register HEIF opener | |
| pillow_heif.register_heif_opener() | |
| LOG = logging.getLogger(__name__) | |
| class ImageBundle: | |
| def __init__(self, image_source=None, image_array=None): | |
| """ | |
| Initialize the ImageHandler object. | |
| :param image_source: Path to the image file or URL of the image. | |
| :param image_array: Numpy array of the image. | |
| """ | |
| self.image_source = image_source | |
| self.image_array = image_array | |
| self.exif_data = {} | |
| self.segmentation_maps = {} | |
| if image_array is not None: | |
| self._open_image_from_array(image_array) | |
| else: | |
| self._open_image() | |
| def _open_image(self): | |
| """ | |
| Open the image and preserve EXIF data. | |
| """ | |
| try: | |
| if is_url(self.image_source): | |
| self._open_image_from_url(self.image_source) | |
| else: | |
| self._open_image_from_path(self.image_source) | |
| self._handle_color_profile() | |
| self._extract_exif_data() | |
| self._rotate_image_if_needed() | |
| self.is_black_and_white = self._is_black_and_white() | |
| self.basename, self.ext = extract_filename_and_extension(self.image_source) | |
| if self.is_black_and_white: | |
| raise ValueError( | |
| "The image is black and white. Please provide a colored image." | |
| ) | |
| except Exception as e: | |
| raise Exception("Corrupt image") | |
| def _open_image_from_path(self, image_path): | |
| """ | |
| Open an image from a file path. | |
| :param image_path: Path to the image file. | |
| """ | |
| self.image = Image.open(image_path) | |
| def _open_image_from_url(self, image_url): | |
| """ | |
| Open an image from a URL. | |
| :param image_url: URL of the image. | |
| """ | |
| response = requests.get(image_url) | |
| response.raise_for_status() # Raise an exception for HTTP errors | |
| self.image = Image.open(BytesIO(response.content)) | |
| def _open_image_from_array(self, image_array): | |
| """ | |
| Open an image from a numpy array. | |
| :param image_array: Numpy array of the image. | |
| """ | |
| self.image = Image.fromarray(image_array) | |
| self._handle_color_profile() | |
| self._extract_exif_data() | |
| self.is_black_and_white = self._is_black_and_white() | |
| self.basename = "uploaded_image" | |
| self.ext = ".jpg" | |
| if self.is_black_and_white: | |
| raise ValueError( | |
| "The image is black and white. Please provide a colored image." | |
| ) | |
| def _handle_color_profile(self): | |
| """ | |
| Handle the color profile of the image if mentioned. | |
| """ | |
| if "icc_profile" in self.image.info: # type: ignore | |
| icc_profile = self.image.info.get("icc_profile") # type: ignore | |
| if icc_profile: | |
| io = BytesIO(icc_profile) | |
| src_profile = ImageCms.ImageCmsProfile(io) | |
| dst_profile = ImageCms.createProfile("sRGB") | |
| self.image = ImageCms.profileToProfile( | |
| self.image, src_profile, dst_profile # type: ignore | |
| ) | |
| def _extract_exif_data(self): | |
| """ | |
| Extract EXIF data from the image. | |
| """ | |
| if hasattr(self.image, "_getexif"): | |
| exif_info = self.image._getexif() # type: ignore | |
| if exif_info is not None: | |
| for tag, value in exif_info.items(): | |
| decoded_tag = ExifTags.TAGS.get(tag, tag) | |
| self.exif_data[decoded_tag] = value | |
| def _rotate_image_if_needed(self): | |
| """ | |
| Rotate the image based on EXIF orientation information. | |
| """ | |
| if not self.image.info: | |
| return | |
| for orientation in ExifTags.TAGS.keys(): | |
| if ExifTags.TAGS[orientation] == "Orientation": | |
| break | |
| exif = dict(self.image.info.items()) | |
| orientation = exif.get(orientation) | |
| if orientation == 3: | |
| self.image = self.image.rotate(180, expand=True) | |
| elif orientation == 6: | |
| self.image = self.image.rotate(270, expand=True) | |
| elif orientation == 8: | |
| self.image = self.image.rotate(90, expand=True) | |
| def _is_black_and_white(self): | |
| """ | |
| Check if the image is black and white even if it has 3 channels. | |
| :return: True if the image is black and white, otherwise False. | |
| """ | |
| if self.image.mode not in ["RGB", "RGBA"]: | |
| return self.image.mode in ["1", "L"] | |
| np_image = np.array(self.image) | |
| if len(np_image.shape) == 3 and np_image.shape[2] == 3: | |
| r, g, b = np_image[:, :, 0], np_image[:, :, 1], np_image[:, :, 2] | |
| return np.all(r == g) and np.all(g == b) | |
| return False | |
| def detect_faces_and_landmarks(self): | |
| """ | |
| Detect faces and landmarks using MediaPipe. | |
| :return: List of dictionaries with face and landmark information. | |
| """ | |
| image_np = self.numpy_image() | |
| face_data = detect_faces_and_landmarks(image_np) | |
| self.faces = face_data | |
| return self.faces | |
| def detect_face_landmarks(self): | |
| """ | |
| Detect face landmarks using MediaPipe. | |
| :return: Dictionary with landmarks for iris, lips, eyebrows, and eyes. | |
| """ | |
| image_np = self.numpy_image() | |
| landmarks = detect_face_landmarks(image_np) | |
| self.landmarks = landmarks | |
| return self.landmarks | |
| def segment_image(self): | |
| """ | |
| Segment the image using MediaPipe Multi-Class Selfie Segmentation. | |
| :return: Dictionary of segmentation masks. | |
| """ | |
| image_np = self.numpy_image() | |
| # save image to check bgr or rgb | |
| cv2.imwrite("workspace/image.jpg", image_np) | |
| masks = mediapipe_selfie_segmentor(image_np) | |
| # Detect face landmarks and create masks for individual features | |
| landmarks = self.detect_face_landmarks() | |
| feature_masks = create_feature_masks(image_np, landmarks) | |
| # Subtract feature masks from face skin mask | |
| for feature in [ | |
| "lips_mask", | |
| "left_eyebrow_mask", | |
| "right_eyebrow_mask", | |
| "left_eye_mask", | |
| "right_eye_mask", | |
| "left_iris_mask", | |
| "right_iris_mask", | |
| ]: | |
| if "iris" in feature: | |
| masks[feature] = feature_masks[feature] | |
| masks["face_skin_mask"] = cv2.subtract( | |
| masks["face_skin_mask"], feature_masks[feature] | |
| ) | |
| self.segmentation_maps = masks | |
| return self.segmentation_maps | |
| def numpy_image(self): | |
| """ | |
| Convert the image to a numpy array. | |
| :return: Numpy array of the image. | |
| """ | |
| image = np.array(self.image) | |
| if image.shape[2] == 4: | |
| image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) | |
| # import IPython; IPython.embed() | |
| # image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| return image | |