Начал менять логику pipeline

src/interpolator.py

@@ -4,7 +4,6 @@ from pathlib import Path
 import torch
 import numpy as np
 from omegaconf import OmegaConf, DictConfig
-from imageio import imread, imwrite
 
 from src.utils.torch import img2tensor, check_dim_and_resize, tensor2img
 from src.utils.build import build_from_cfg
@@ -83,7 +82,7 @@ class ImageInterpolator:
             f"Initialized ImageInterpolator with device: {device}, anchor: {anchor}, available VRAM: {self.vram_available} bytes"
         )
 
-    def interpolate(self, image1: Path, image2: Path, output_path: Path):
+    def interpolate(self, image1: np.ndarray, image2: np.ndarray) -> np.ndarray:
         """
         Interpolates between two images and saves the result.
         Args:
@@ -92,8 +91,8 @@ class ImageInterpolator:
             output_path (Path): Path to save the interpolated image (only png and jpg formats are supported)
         """
         logging.debug(f"Reading images: {image1} and {image2}")
-        tensor1 = img2tensor(imread(image1)).to(self.device)
-        tensor2 = img2tensor(imread(image2)).to(self.device)
+        tensor1 = img2tensor(image1).to(self.device)
+        tensor2 = img2tensor(image2).to(self.device)
         logging.debug(
             f"Image shapes after conversion to tensors: {tensor1.shape}, {tensor2.shape}"
         )
@@ -122,8 +121,7 @@ class ImageInterpolator:
         logging.debug(f"Interpolated image shape before unpadding: {interpolated.shape}")
         (interpolated,) = padder.unpad(interpolated)
         logging.debug(f"Interpolated image shape after unpadding: {interpolated.shape}")
-        imwrite(output_path, tensor2img(interpolated.cpu()))
-        logging.debug(f"Saved interpolated image to: {output_path}")
+        return tensor2img(interpolated.cpu())
 
     def scale(self, height: int, width: int) -> float:
         scale = (