# 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
        # This is torch format rather than PIL image
        self.image_shape = (self.imsize, self.imsize)
        self.image_length = self.channels * self.imsize * self.imsize
        img_space = gym.spaces.Box(
            0, 1, (self.image_length,), dtype=np.float32)
        self._image_goal = img_space.sample()  # has to be done for presampling
        spaces = self.wrapped_env.observation_space.spaces.copy()
        spaces['observation'] = img_space
        spaces['desired_goal'] = img_space
        spaces['achieved_goal'] = img_space
        spaces['image_observation'] = img_space
        spaces['image_desired_goal'] = img_space
        spaces['image_achieved_goal'] = img_space
        self.observation_space = gym.spaces.Dict(spaces)

        self.action_space = self.wrapped_env.action_space
        self.reward_type = reward_type
        self.threshold = threshold
        self._last_image = None

        self._presampled_goals = None
        if presampled_goals is not None:
            assert 'image_desired_goal' in presampled_goals, presampled_goals.keys()

            # Normalize goal images.
            image_goals = presampled_goals['image_desired_goal'].copy()
            image_goals_max = np.max(image_goals)
            if image_goals_max > 1:
                assert image_goals_max <= 255, image_goals_max
                assert np.min(image_goals) >= 0, np.min(image_goals)

# 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

        # This is torch format rather than PIL image
        self.image_shape = (self.imsize, self.imsize)
        self.image_length = self.channels * self.imsize * self.imsize
        img_space = gym.spaces.Box(
            0, 1, (self.image_length,), dtype=np.float32)
        self._image_goal = img_space.sample()  # has to be done for presampling
        spaces = self.wrapped_env.observation_space.spaces.copy()
        spaces['observation'] = img_space
        spaces['desired_goal'] = img_space
        spaces['achieved_goal'] = img_space
        spaces['image_observation'] = img_space
        spaces['image_desired_goal'] = img_space
        spaces['image_achieved_goal'] = img_space
        self.observation_space = gym.spaces.Dict(spaces)

        self.action_space = self.wrapped_env.action_space
        self.reward_type = reward_type
        self.threshold = threshold
        self._last_image = None

        self._presampled_goals = None
        if presampled_goals is not None:
            assert 'image_desired_goal' in presampled_goals, presampled_goals.keys()

            # Normalize goal images.
            image_goals = presampled_goals['image_desired_goal'].copy()
            image_goals_max = np.max(image_goals)
            if image_goals_max > 1:
                assert image_goals_max <= 255, image_goals_max
                assert np.min(image_goals) >= 0, np.min(image_goals)

                image_goals = normalize_image(image_goals)