switch info to debug in logger
This commit is contained in:
31
main.py
31
main.py
@@ -84,44 +84,44 @@ class ImageInterpolator:
|
|||||||
self.vram_available = get_vram_available(device)
|
self.vram_available = get_vram_available(device)
|
||||||
self.embt = torch.tensor(1 / 2).float().view(1, 1, 1, 1).to(device)
|
self.embt = torch.tensor(1 / 2).float().view(1, 1, 1, 1).to(device)
|
||||||
self.model_runner = model_runner
|
self.model_runner = model_runner
|
||||||
logging.info(
|
logging.debug(
|
||||||
f"Initialized ImageInterpolator with device: {device}, anchor: {anchor}, available VRAM: {self.vram_available} bytes"
|
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: Path, image2: Path, output_path: Path):
|
||||||
logging.info(f"Reading images: {image1} and {image2}")
|
logging.debug(f"Reading images: {image1} and {image2}")
|
||||||
tensor1 = img2tensor(utils.read(image1)).to(self.device)
|
tensor1 = img2tensor(utils.read(image1)).to(self.device)
|
||||||
tensor2 = img2tensor(utils.read(image2)).to(self.device)
|
tensor2 = img2tensor(utils.read(image2)).to(self.device)
|
||||||
logging.info(
|
logging.debug(
|
||||||
f"Image shapes after conversion to tensors: {tensor1.shape}, {tensor2.shape}"
|
f"Image shapes after conversion to tensors: {tensor1.shape}, {tensor2.shape}"
|
||||||
)
|
)
|
||||||
tensor1, tensor2 = check_dim_and_resize(tensor1, tensor2)
|
tensor1, tensor2 = check_dim_and_resize(tensor1, tensor2)
|
||||||
logging.info(f"Image shapes after resizing: {tensor1.shape}, {tensor2.shape}")
|
logging.debug(f"Image shapes after resizing: {tensor1.shape}, {tensor2.shape}")
|
||||||
h, w = tensor1.shape[2], tensor1.shape[3]
|
h, w = tensor1.shape[2], tensor1.shape[3]
|
||||||
logging.info(f"Interpolating images of size: {h}x{w}")
|
logging.debug(f"Interpolating images of size: {h}x{w}")
|
||||||
|
|
||||||
scale = self.scale(h, w)
|
scale = self.scale(h, w)
|
||||||
logging.info(f"Calculated scale factor: {scale:.2f}")
|
logging.debug(f"Calculated scale factor: {scale:.2f}")
|
||||||
padding = int(16 / scale)
|
padding = int(16 / scale)
|
||||||
logging.info(f"Calculated padding: {padding} pixels")
|
logging.debug(f"Calculated padding: {padding} pixels")
|
||||||
padder = InputPadder(tensor1.shape, divisor=padding)
|
padder = InputPadder(tensor1.shape, divisor=padding)
|
||||||
tensor1_padded, tensor2_padded = padder.pad(tensor1, tensor2)
|
tensor1_padded, tensor2_padded = padder.pad(tensor1, tensor2)
|
||||||
logging.info(
|
logging.debug(
|
||||||
f"Image shapes after padding: {tensor1_padded.shape}, {tensor2_padded.shape}"
|
f"Image shapes after padding: {tensor1_padded.shape}, {tensor2_padded.shape}"
|
||||||
)
|
)
|
||||||
|
|
||||||
tensor1_padded = tensor1_padded.to(self.device)
|
tensor1_padded = tensor1_padded.to(self.device)
|
||||||
tensor2_padded = tensor2_padded.to(self.device)
|
tensor2_padded = tensor2_padded.to(self.device)
|
||||||
logging.info("Running model inference for interpolation")
|
logging.debug("Running model inference for interpolation")
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
interpolated = self.model_runner.model(
|
interpolated = self.model_runner.model(
|
||||||
tensor1_padded, tensor2_padded, self.embt, scale_factor=scale, eval=True
|
tensor1_padded, tensor2_padded, self.embt, scale_factor=scale, eval=True
|
||||||
)["imgt_pred"]
|
)["imgt_pred"]
|
||||||
logging.info(f"Interpolated image shape before unpadding: {interpolated.shape}")
|
logging.debug(f"Interpolated image shape before unpadding: {interpolated.shape}")
|
||||||
(interpolated,) = padder.unpad(interpolated)
|
(interpolated,) = padder.unpad(interpolated)
|
||||||
logging.info(f"Interpolated image shape after unpadding: {interpolated.shape}")
|
logging.debug(f"Interpolated image shape after unpadding: {interpolated.shape}")
|
||||||
utils.write(output_path, tensor2img(interpolated.cpu()))
|
utils.write(output_path, tensor2img(interpolated.cpu()))
|
||||||
logging.info(f"Saved interpolated image to: {output_path}")
|
logging.debug(f"Saved interpolated image to: {output_path}")
|
||||||
|
|
||||||
def scale(self, height: int, width: int) -> float:
|
def scale(self, height: int, width: int) -> float:
|
||||||
scale = (
|
scale = (
|
||||||
@@ -145,7 +145,9 @@ def main():
|
|||||||
ckpt_path = Path("src/pretrained/amt-g.pth")
|
ckpt_path = Path("src/pretrained/amt-g.pth")
|
||||||
image1_path = Path("source/img0.png")
|
image1_path = Path("source/img0.png")
|
||||||
image2_path = Path("source/img1.png")
|
image2_path = Path("source/img1.png")
|
||||||
output_path = Path("output/interpolated_image.png")
|
image3_path = Path("source/img2.png")
|
||||||
|
output_path1 = Path("output/interpolated_image1.png")
|
||||||
|
output_path2 = Path("output/interpolated_image2.png")
|
||||||
|
|
||||||
device = get_device()
|
device = get_device()
|
||||||
model_runner = ModelRunner(config_path, ckpt_path, device)
|
model_runner = ModelRunner(config_path, ckpt_path, device)
|
||||||
@@ -166,7 +168,8 @@ def main():
|
|||||||
else:
|
else:
|
||||||
raise Exception(f"Unsupported device type: {device.type}")
|
raise Exception(f"Unsupported device type: {device.type}")
|
||||||
interpolator = ImageInterpolator(device, anchor, model_runner)
|
interpolator = ImageInterpolator(device, anchor, model_runner)
|
||||||
interpolator.interpolate(image1_path, image2_path, output_path)
|
interpolator.interpolate(image1_path, image2_path, output_path1)
|
||||||
|
interpolator.interpolate(image2_path, image3_path, output_path2)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Reference in New Issue
Block a user