# Path: weakly_supervised_control/envs/env_util.py
# def normalize_image(image: np.ndarray, dtype=np.float64) -> np.ndarray:
#     assert image.dtype == np.uint8
#     return dtype(image) / 255.0
# 
# def unormalize_image(image: np.ndarray) -> np.ndarray:
#     assert image.dtype != np.uint8
#     return np.uint8(image * 255.0)
# 
# def concat_images(images: np.ndarray,
#                   resize_shape: Tuple[int, int] = None,
#                   concat_horizontal: bool = True):
#     """Concatenates N images into a single image.
#     Args:
#       images: An RGB array of shape (N, width, height, 3)
#       resize_shape: (width, height)
#       concat_horizontal: If true, returns an image of size (N * width, height).
#                          Otherwise returns an image of size (width, N * height).
#     """
#     assert len(images.shape) == 4, images.shape
#     assert images.shape[3] == 3, images.shape
#     if resize_shape:
#         images = np.array([cv2.resize(image, resize_shape)
#                            for image in images])
#     if concat_horizontal:
#         stacked_images = np.concatenate(images, axis=1)
#     else:
#         stacked_images = np.concatenate(images, axis=0)
#     return stacked_images


from typing import Dict, List, Tuple
from collections import OrderedDict
from numbers import Number
from PIL import Image
from multiworld.core.multitask_env import MultitaskEnv
from multiworld.core.wrapper_env import ProxyEnv
from multiworld.envs.env_util import concatenate_box_spaces
from weakly_supervised_control.envs.env_util import normalize_image, unormalize_image, concat_images
import random
import warnings
import cv2
import gym
import numpy as np
                assert np.min(image_goals) >= 0, np.min(image_goals)
                image_goals = normalize_image(image_goals)

            if len(image_goals[0].shape) != 1:
                # Flatten images.
                presampled_goals['image_desired_goal'] = np.array(
                    [self._flatten_image(image) for image in image_goals])
            self._presampled_goals = presampled_goals
            self.num_goals_presampled = presampled_goals['image_desired_goal'].shape[0]

    def _flatten_image(self, image: np.ndarray) -> np.ndarray:
        # Flatten the image from (w, h, 3n) to (3n * w * h).
        image = image.transpose()
        assert image.shape[0] == self.channels, image.shape
        return image.flatten()
    
    def  _unflatten_image(self, image: np.ndarray) -> np.ndarray:
        image = image.reshape(
                self.channels,
                self.imsize,
                self.imsize,
            ).transpose()
        assert image.shape[2] == self.channels,  image.shape
        return image

    def obs_to_image(self, image_obs: np.ndarray) -> np.ndarray:
        # Input obs is of shape (3n * w * h)
        image_obs = self._unflatten_image(image_obs)  # (w, h, 3n)
        images = np.split(image_obs, self.channels / 3, axis=2)  # n images of shape (w, h, 3)
        images = np.concatenate([np.expand_dims(img, 0) for img in images])  # (n, w, h, 3)

# Path: weakly_supervised_control/envs/env_util.py
# def normalize_image(image: np.ndarray, dtype=np.float64) -> np.ndarray:
#     assert image.dtype == np.uint8
#     return dtype(image) / 255.0
# 
# def unormalize_image(image: np.ndarray) -> np.ndarray:
#     assert image.dtype != np.uint8
#     return np.uint8(image * 255.0)
# 
# def concat_images(images: np.ndarray,
#                   resize_shape: Tuple[int, int] = None,
#                   concat_horizontal: bool = True):
#     """Concatenates N images into a single image.
#     Args:
#       images: An RGB array of shape (N, width, height, 3)
#       resize_shape: (width, height)
#       concat_horizontal: If true, returns an image of size (N * width, height).
#                          Otherwise returns an image of size (width, N * height).
#     """
#     assert len(images.shape) == 4, images.shape
#     assert images.shape[3] == 3, images.shape
#     if resize_shape:
#         images = np.array([cv2.resize(image, resize_shape)
#                            for image in images])
#     if concat_horizontal:
#         stacked_images = np.concatenate(images, axis=1)
#     else:
#         stacked_images = np.concatenate(images, axis=0)
#     return stacked_images


# Path: weakly_supervised_control/envs/image_env.py
from typing import Dict, List, Tuple
from collections import OrderedDict
from numbers import Number
from PIL import Image
from multiworld.core.multitask_env import MultitaskEnv
from multiworld.core.wrapper_env import ProxyEnv
from multiworld.envs.env_util import concatenate_box_spaces
from weakly_supervised_control.envs.env_util import normalize_image, unormalize_image, concat_images
import random
import warnings
import cv2
import gym
import numpy as np

                assert np.min(image_goals) >= 0, np.min(image_goals)
                image_goals = normalize_image(image_goals)

            if len(image_goals[0].shape) != 1:
                # Flatten images.
                presampled_goals['image_desired_goal'] = np.array(
                    [self._flatten_image(image) for image in image_goals])
            self._presampled_goals = presampled_goals
            self.num_goals_presampled = presampled_goals['image_desired_goal'].shape[0]

    def _flatten_image(self, image: np.ndarray) -> np.ndarray:
        # Flatten the image from (w, h, 3n) to (3n * w * h).
        image = image.transpose()
        assert image.shape[0] == self.channels, image.shape
        return image.flatten()
    
    def  _unflatten_image(self, image: np.ndarray) -> np.ndarray:
        image = image.reshape(
                self.channels,
                self.imsize,
                self.imsize,
            ).transpose()
        assert image.shape[2] == self.channels,  image.shape
        return image

    def obs_to_image(self, image_obs: np.ndarray) -> np.ndarray:
        # Input obs is of shape (3n * w * h)
        image_obs = self._unflatten_image(image_obs)  # (w, h, 3n)
        images = np.split(image_obs, self.channels / 3, axis=2)  # n images of shape (w, h, 3)
        images = np.concatenate([np.expand_dims(img, 0) for img in images])  # (n, w, h, 3)

        stacked_images = concat_images(images)  # (n * w, h, 3)