Merge remote-tracking branch 'origin/feat-pool-encoder' into feat-poo…

…l-encoder
gnes-ai · Aug 26, 2019 · 3600c29 · 3600c29
2 parents ce0e65a + 615851e
commit 3600c29
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 10 deletions.
diff --git a/gnes/encoder/image/torchvision.py b/gnes/encoder/image/torchvision.py
@@ -109,14 +109,9 @@ def _encode(_, img: List['np.ndarray']):
             if self._use_cuda:
                 img_tensor = img_tensor.cuda()
 
-            result_npy = []
-            for t in img_tensor:
-                t = torch.unsqueeze(t, 0)
-                encodes = self._model(t)
-                encodes = torch.squeeze(encodes, 0)
-                result_npy.append(encodes.data.cpu().numpy())
-
-            output = np.array(result_npy, dtype=np.float32)
+            encodes = self._model(img_tensor)
+
+            output = np.array(encodes.data.cpu().numpy(), dtype=np.float32)
             return output
 
         output = _encode(self, img)

diff --git a/gnes/preprocessor/helper.py b/gnes/preprocessor/helper.py
@@ -135,7 +135,7 @@ def split_video_frames(buffer_data: bytes,
     return [np.array(Image.open(io.BytesIO(chunk))) for chunk in chunks]
 
 
-def get_gif(images, fps=4):
+def get_gif(images: 'np.ndarray', fps=10):
     cmd = ['ffmpeg', '-y',
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',

diff --git a/gnes/preprocessor/io_utils/__init__.py b/gnes/preprocessor/io_utils/__init__.py
diff --git a/gnes/preprocessor/io_utils/gif.py b/gnes/preprocessor/io_utils/gif.py
@@ -47,7 +47,7 @@ def decode_gif(data: bytes, fps: int = -1,
 
 
 def encode_gif(
-        images: List[np.ndarray],
+        images: np.ndarray,
         scale: str,
         fps: int,
         pix_fmt: str = 'rgb24'):