fix: vsr vs_inference (#10)

TensoRaws · Oct 9, 2024 · cee4763 · cee4763
1 parent ef11d44
commit cee4763
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/ccrestoration/model/vsr_base_model.py b/ccrestoration/model/vsr_base_model.py
@@ -51,7 +51,7 @@ def inference_image_list(self, img_list: List[np.ndarray]) -> List[np.ndarray]:
 
             for i in range(out.shape[1]):
                 img = out[0, i, :, :, :]
-                img = img.squeeze(0).permute(1, 2, 0).cpu().numpy()
+                img = img.permute(1, 2, 0).cpu().numpy()
                 img = (img * 255).clip(0, 255).astype("uint8")
                 img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                 res_img_list.append(img)

diff --git a/ccrestoration/vs/vsr.py b/ccrestoration/vs/vsr.py
@@ -59,7 +59,7 @@ def _inference(n: int, f: list[vs.VideoFrame]) -> vs.VideoFrame:
             output = inference(img)
 
             for i in range(output.shape[0]):
-                cache[n + i] = output[i, :, :, :]
+                cache[n + i] = output[0, i, :, :, :]
 
         return _tensor_to_frame(cache[n], f[1].copy())
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -40,7 +40,7 @@ license = "MIT"
 name = "ccrestoration"
 readme = "README.md"
 repository = "https://github.com/TensoRaws/ccrestoration"
-version = "0.0.6"
+version = "0.0.7"
 
 # Requirements
 [tool.poetry.dependencies]