Merge branch 'main'

Dreamer/dmc2gym/wrappers.py

@@ -122,7 +122,7 @@ class DMCWrapper(core.Env):
                 if img_source == "images":
                     self._bg_source = natural_imgsource.RandomImageSource(shape2d, files, grayscale=False,  max_videos=100, random_bg=False)
                 elif img_source == "video":
-                    self._bg_source = natural_imgsource.RandomVideoSource(shape2d, files, grayscale=False,max_videos=100, random_bg=False)
+                    self._bg_source = natural_imgsource.RandomVideoSource(shape2d, files, grayscale=False, max_videos=100, random_bg=False)
                 else:
                     raise Exception("img_source %s not defined." % img_source)
 

Dreamer/dreamers.py

@@ -308,6 +308,7 @@ class SeparationDreamer(Dreamer):
         with tf.GradientTape(persistent=True) as model_tape:
 
             # main
+            data["image"] = tf.transpose(data["image"], perm=[0, 1, 3, 4, 2])
             embed = self._encode(data)
             post, prior = self._dynamics.observe(embed, data['action'])
             feat = self._dynamics.get_feat(post)

Dreamer/tools.py

@@ -87,24 +87,26 @@ def video_summary(name, video, step=None, fps=20):
 def encode_gif(frames, fps):
     from subprocess import Popen, PIPE
     print(frames[0].shape)
-    if frames[0].shape[-1] != 3:
+    if frames[0].shape[-1] > 3:
         frames =  np.transpose(frames, [0, 2, 3, 1])
     h, w, c = frames[0].shape
-    print(h,w,c)
+    print(frames[0].shape)
-    pxfmt = {1: 'gray', 3: 'rgb24'}[c]
+
-    cmd = ' '.join([
+    if c!=64:
-        f'ffmpeg -y -f rawvideo -vcodec rawvideo',
+        pxfmt = {1: 'gray', 3: 'rgb24'}[c]
-        f'-r {fps:.02f} -s {w}x{h} -pix_fmt {pxfmt} -i - -filter_complex',
+        cmd = ' '.join([
-        f'[0:v]split[x][z];[z]palettegen[y];[x]fifo[x];[x][y]paletteuse',
+            f'ffmpeg -y -f rawvideo -vcodec rawvideo',
-        f'-r {fps:.02f} -f gif -'])
+            f'-r {fps:.02f} -s {w}x{h} -pix_fmt {pxfmt} -i - -filter_complex',
-    proc = Popen(cmd.split(' '), stdin=PIPE, stdout=PIPE, stderr=PIPE)
+            f'[0:v]split[x][z];[z]palettegen[y];[x]fifo[x];[x][y]paletteuse',
-    for image in frames:
+            f'-r {fps:.02f} -f gif -'])
-        proc.stdin.write(image.tostring())
+        proc = Popen(cmd.split(' '), stdin=PIPE, stdout=PIPE, stderr=PIPE)
-    out, err = proc.communicate()
+        for image in frames:
-    if proc.returncode:
+            proc.stdin.write(image.tostring())
-        raise IOError('\n'.join([' '.join(cmd), err.decode('utf8')]))
+        out, err = proc.communicate()
-    del proc
+        if proc.returncode:
-    return out
+            raise IOError('\n'.join([' '.join(cmd), err.decode('utf8')]))
+        del proc
+        return out
 
 
 def simulate(agent, envs, steps=0, episodes=0, state=None):

Dreamer/train_configs/dreamer.yaml

@@ -1,6 +1,6 @@
 dmc: 
 
-  logdir: /home/vedant/tia/Dreamer/logdir
+  logdir: /media/vedant/cpsDataStorageWK/Vedant/tia_logs
   video_dir_train: /media/vedant/cpsDataStorageWK/Vedant/natural_video_setting/train/
   video_dir_test: /media/vedant/cpsDataStorageWK/Vedant/natural_video_setting/test/
   debug: False

Dreamer/train_configs/tia.yaml

@@ -1,6 +1,6 @@
 dmc:
 
-  logdir: /home/vedant/tia/Dreamer/logdir
+  logdir: /media/vedant/cpsDataStorageWK/Vedant/tia_logs
   video_dir_train: /media/vedant/cpsDataStorageWK/Vedant/natural_video_setting/train/
   video_dir_test: /media/vedant/cpsDataStorageWK/Vedant/natural_video_setting/test/
   debug: False