Adding files

This commit is contained in:
ved1 2023-06-01 16:23:18 +02:00
parent 4d57be91cd
commit 6c82d972f8
23 changed files with 54 additions and 146 deletions

View File

@ -1,36 +1,19 @@
# This code provides the class that is used to generate backgrounds for the natural background setting
# the class is used inside an environment wrapper and will be called each time the env generates an observation
# the code is largely based on https://github.com/facebookresearch/deep_bisim4control
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import numpy as np
import cv2
import skvideo.io
import random import random
import tqdm
class BackgroundMatting(object): import cv2
""" import numpy as np
Produce a mask by masking the given color. This is a simple strategy import skvideo.io
but effective for many games.
"""
def __init__(self, color):
"""
Args:
color: a (r, g, b) tuple or single value for grayscale
"""
self._color = color
def get_mask(self, img):
return img == self._color
class ImageSource(object): class ImageSource(object):
""" """
Source of natural images to be added to a simulated environment. Source of natural images to be added to a simulated environment.
""" """
def get_image(self): def get_image(self):
""" """
Returns: Returns:
@ -43,141 +26,56 @@ class ImageSource(object):
pass pass
class FixedColorSource(ImageSource):
def __init__(self, shape, color):
"""
Args:
shape: [h, w]
color: a 3-tuple
"""
self.arr = np.zeros((shape[0], shape[1], 3))
self.arr[:, :] = color
def get_image(self):
return self.arr
class RandomColorSource(ImageSource):
def __init__(self, shape):
"""
Args:
shape: [h, w]
"""
self.shape = shape
self.arr = None
self.reset()
def reset(self):
self._color = np.random.randint(0, 256, size=(3,))
self.arr = np.zeros((self.shape[0], self.shape[1], 3))
self.arr[:, :] = self._color
def get_image(self):
return self.arr
class NoiseSource(ImageSource):
def __init__(self, shape, strength=255):
"""
Args:
shape: [h, w]
strength (int): the strength of noise, in range [0, 255]
"""
self.shape = shape
self.strength = strength
def get_image(self):
return np.random.randn(self.shape[0], self.shape[1], 3) * self.strength
class RandomImageSource(ImageSource):
def __init__(self, shape, filelist, total_frames=None, grayscale=False):
"""
Args:
shape: [h, w]
filelist: a list of image files
"""
self.grayscale = grayscale
self.total_frames = total_frames
self.shape = shape
self.filelist = filelist
self.build_arr()
self.current_idx = 0
self.reset()
def build_arr(self):
self.total_frames = self.total_frames if self.total_frames else len(self.filelist)
self.arr = np.zeros((self.total_frames, self.shape[0], self.shape[1]) + ((3,) if not self.grayscale else (1,)))
for i in range(self.total_frames):
# if i % len(self.filelist) == 0: random.shuffle(self.filelist)
fname = self.filelist[i % len(self.filelist)]
if self.grayscale: im = cv2.imread(fname, cv2.IMREAD_GRAYSCALE)[..., None]
else: im = cv2.imread(fname, cv2.IMREAD_COLOR)
self.arr[i] = cv2.resize(im, (self.shape[1], self.shape[0])) ## THIS IS NOT A BUG! cv2 uses (width, height)
def reset(self):
self._loc = np.random.randint(0, self.total_frames)
def get_image(self):
return self.arr[self._loc]
class RandomVideoSource(ImageSource): class RandomVideoSource(ImageSource):
def __init__(self, shape, filelist, total_frames=None, grayscale=False): def __init__(self, shape, filelist, random_bg=False, max_videos=50, grayscale=False):
""" """
Args: Args:
shape: [h, w] shape: [h, w]
filelist: a list of video files filelist: a list of video files
""" """
self.grayscale = grayscale self.grayscale = grayscale
self.total_frames = total_frames
self.shape = shape self.shape = shape
self.filelist = filelist self.filelist = filelist
self.build_arr() random.shuffle(self.filelist)
self.filelist = self.filelist[:max_videos]
self.max_videos = max_videos
self.random_bg = random_bg
self.current_idx = 0 self.current_idx = 0
self._current_vid = None
self.reset() self.reset()
def build_arr(self): def load_video(self, vid_id):
if not self.total_frames: fname = self.filelist[vid_id]
self.total_frames = 0 if self.grayscale:
self.arr = None frames = skvideo.io.vread(fname, outputdict={"-pix_fmt": "gray"})
random.shuffle(self.filelist)
for fname in tqdm.tqdm(self.filelist, desc="Loading videos for natural", position=0):
if self.grayscale: frames = skvideo.io.vread(fname, outputdict={"-pix_fmt": "gray"})
else: frames = skvideo.io.vread(fname)
local_arr = np.zeros((frames.shape[0], self.shape[0], self.shape[1]) + ((3,) if not self.grayscale else (1,)))
for i in tqdm.tqdm(range(frames.shape[0]), desc="video frames", position=1):
local_arr[i] = cv2.resize(frames[i], (self.shape[1], self.shape[0])) ## THIS IS NOT A BUG! cv2 uses (width, height)
if self.arr is None:
self.arr = local_arr
else:
self.arr = np.concatenate([self.arr, local_arr], 0)
self.total_frames += local_arr.shape[0]
else: else:
self.arr = np.zeros((self.total_frames, self.shape[0], self.shape[1]) + ((3,) if not self.grayscale else (1,))) frames = skvideo.io.vread(fname, num_frames=1000)
total_frame_i = 0
file_i = 0
with tqdm.tqdm(total=self.total_frames, desc="Loading videos for natural") as pbar:
while total_frame_i < self.total_frames:
if file_i % len(self.filelist) == 0: random.shuffle(self.filelist)
file_i += 1
fname = self.filelist[file_i % len(self.filelist)]
if self.grayscale: frames = skvideo.io.vread(fname, outputdict={"-pix_fmt": "gray"})
else: frames = skvideo.io.vread(fname)
for frame_i in range(frames.shape[0]):
if total_frame_i >= self.total_frames: break
if self.grayscale:
self.arr[total_frame_i] = cv2.resize(frames[frame_i], (self.shape[1], self.shape[0]))[..., None] ## THIS IS NOT A BUG! cv2 uses (width, height)
else:
self.arr[total_frame_i] = cv2.resize(frames[frame_i], (self.shape[1], self.shape[0]))
pbar.update(1)
total_frame_i += 1
img_arr = np.zeros((frames.shape[0], self.shape[0], self.shape[1]) + ((3,) if not self.grayscale else (1,)))
for i in range(frames.shape[0]):
if self.grayscale:
img_arr[i] = cv2.resize(frames[i], (self.shape[1], self.shape[0]))[..., None] # THIS IS NOT A BUG! cv2 uses (width, height)
else:
img_arr[i] = cv2.resize(frames[i], (self.shape[1], self.shape[0]))
return img_arr
def reset(self): def reset(self):
self._loc = np.random.randint(0, self.total_frames) del self._current_vid
self._video_id = np.random.randint(0, len(self.filelist))
self._current_vid = self.load_video(self._video_id)
while True:
try:
self._video_id = np.random.randint(0, len(self.filelist))
self._current_vid = self.load_video(self._video_id)
break
except Exception:
continue
self._loc = np.random.randint(0, len(self._current_vid))
def get_image(self): def get_image(self):
img = self.arr[self._loc % self.total_frames] if self.random_bg:
self._loc += 1 self._loc = np.random.randint(0, len(self._current_vid))
else:
self._loc += 1
img = self._current_vid[self._loc % len(self._current_vid)]
return img return img

View File

@ -8,6 +8,11 @@ import skimage.io
from dmc2gym import natural_imgsource from dmc2gym import natural_imgsource
high_noise = False
def set_global_var(set_high_noise):
global high_noise
high_noise = set_high_noise
def _spec_to_box(spec): def _spec_to_box(spec):
def extract_min_max(s): def extract_min_max(s):
@ -108,13 +113,16 @@ class DMCWrapper(core.Env):
self._bg_source = natural_imgsource.NoiseSource(shape2d) self._bg_source = natural_imgsource.NoiseSource(shape2d)
else: else:
files = glob.glob(os.path.expanduser(resource_files)) files = glob.glob(os.path.expanduser(resource_files))
self.files = files
self.total_frames = total_frames
self.shape2d = shape2d
assert len(files), "Pattern {} does not match any files".format( assert len(files), "Pattern {} does not match any files".format(
resource_files resource_files
) )
if img_source == "images": if img_source == "images":
self._bg_source = natural_imgsource.RandomImageSource(shape2d, files, grayscale=True, total_frames=total_frames) self._bg_source = natural_imgsource.RandomImageSource(shape2d, files, grayscale=False, max_videos=50, random_bg=False)
elif img_source == "video": elif img_source == "video":
self._bg_source = natural_imgsource.RandomVideoSource(shape2d, files, grayscale=True, total_frames=total_frames) self._bg_source = natural_imgsource.RandomVideoSource(shape2d, files, grayscale=False,max_videos=50, random_bg=False)
else: else:
raise Exception("img_source %s not defined." % img_source) raise Exception("img_source %s not defined." % img_source)
@ -185,6 +193,8 @@ class DMCWrapper(core.Env):
def reset(self): def reset(self):
time_step = self._env.reset() time_step = self._env.reset()
self._bg_source.reset()
#self._bg_source = natural_imgsource.RandomVideoSource(self.shape2d, self.files, grayscale=True, total_frames=self.total_frames, high_noise=high_noise)
obs = self._get_obs(time_step) obs = self._get_obs(time_step)
return obs return obs

Binary file not shown.

Binary file not shown.

Binary file not shown.