442 lines
18 KiB
Python
442 lines
18 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import cv2
|
|
import time
|
|
from PIL import Image
|
|
import numpy as np
|
|
import scipy.ndimage
|
|
|
|
|
|
class MaskExpansionModel(nn.Module):
|
|
"""
|
|
A PyTorch model that expands or erodes a mask using the provided function.
|
|
"""
|
|
|
|
def forward(self, mask, expand, tapered_corners):
|
|
assert mask.dim() == 4, "Input mask must have 3 dimensions (batch_size, height, width)."
|
|
bs = mask.size(0)
|
|
c = 0 if tapered_corners else 1
|
|
kernel = torch.tensor([[c, 1, c],
|
|
[1, 1, 1],
|
|
[c, 1, c]], dtype=mask.dtype, device=mask.device)
|
|
kernel = kernel.view(1, 1, 3, 3)
|
|
|
|
if expand < 0:
|
|
kernel = -kernel
|
|
expand = -expand
|
|
mask = 1 - mask
|
|
|
|
mask = mask.unsqueeze(1)
|
|
# print(mask.shape)
|
|
for _ in range(expand):
|
|
for i in range(bs):
|
|
# print(i)
|
|
mask[i] = torch.nn.functional.conv2d(mask[i], kernel, padding=1)
|
|
mask[i] = torch.clamp(mask[i], 0, 1) # Clip the mask values to the range [0, 1]
|
|
|
|
mask = mask.squeeze(1)
|
|
if expand < 0:
|
|
mask = 1 - mask
|
|
|
|
return mask
|
|
|
|
|
|
class LanczosInterpolationLayer(nn.Module):
|
|
"""
|
|
A PyTorch module that applies Lanczos interpolation to resize a batch of images.
|
|
|
|
Args:
|
|
target_size (tuple): The target size for resizing the images, in the format (width, height).
|
|
"""
|
|
|
|
def __init__(self, target_size):
|
|
super(LanczosInterpolationLayer, self).__init__()
|
|
self.target_size = target_size
|
|
|
|
def visualize_tensor(self,tensor, title):
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
# Convert the tensor to a PIL image
|
|
tensor = tensor.squeeze().cpu().numpy()
|
|
tensor = (tensor * 255).astype('uint8')
|
|
image = Image.fromarray(tensor)
|
|
|
|
# Create a new image for the text
|
|
text_image = Image.new('RGB', (200, 100), color='white')
|
|
draw = ImageDraw.Draw(text_image)
|
|
font = ImageFont.truetype("arial.ttf", 16)
|
|
|
|
# Add the title, min, and max values to the text image
|
|
min_val = tensor.min()
|
|
max_val = tensor.max()
|
|
text = f"{title}\nMin: {min_val:.6f}\nMax: {max_val:.6f}"
|
|
draw.text((10, 10), text, fill='black', font=font)
|
|
|
|
# Combine the original image and the text image
|
|
combined_image = Image.new('RGB', (image.width + text_image.width, max(image.height, text_image.height)),
|
|
color='white')
|
|
combined_image.paste(image, (0, 0))
|
|
combined_image.paste(text_image, (image.width, 0))
|
|
|
|
# Display the combined image
|
|
combined_image.show()
|
|
def forward(self, x):
|
|
"""
|
|
Performs the forward pass of the Lanczos interpolation layer.
|
|
|
|
Args:
|
|
x (torch.Tensor): A batch of input images with shape (batch_size, channels, height, width).
|
|
|
|
Returns:
|
|
torch.Tensor: A batch of resized images with shape (batch_size, channels, target_height, target_width).
|
|
"""
|
|
device = torch.device("cpu")
|
|
interpolated_images = []
|
|
for image in x:
|
|
# Convert the PyTorch tensor to a NumPy array
|
|
image_np = image.permute(1, 2, 0).to("cpu", torch.float32).numpy()
|
|
|
|
# Resize the image using OpenCV's resize function with Lanczos interpolation
|
|
interpolated_image = cv2.resize(image_np, self.target_size, interpolation=cv2.INTER_LANCZOS4)
|
|
mn = interpolated_image.min()
|
|
mx = interpolated_image.max()
|
|
|
|
color_range = mx - mn
|
|
|
|
under_bottom_target_range = -10 // 255
|
|
over_top_target_range = 10 // 255
|
|
|
|
#self.visualize_tensor(torch.from_numpy(interpolated_image), "Original Image")
|
|
|
|
if color_range > 1:
|
|
# here we are going to compress the over the top values
|
|
over_top = interpolated_image > 1
|
|
not_over_top = interpolated_image <= 1
|
|
over_top_range = mx - 1
|
|
over_top_scale = over_top_target_range / over_top_range
|
|
# scale the over the top values to the target range
|
|
interpolated_image[over_top] = 1 + (interpolated_image[over_top] - 1) * over_top_scale
|
|
#self.visualize_tensor(torch.from_numpy(interpolated_image), "After Over-the-Top Compression")
|
|
# scale the non-over-the top to make room
|
|
interpolated_image[not_over_top] = interpolated_image[not_over_top] * (1 - over_top_target_range)
|
|
#self.visualize_tensor(torch.from_numpy(interpolated_image), "After Non-Over-the-Top Scaling")
|
|
|
|
# here we are going to compress the under the bottom values
|
|
under_bottom = interpolated_image < 0
|
|
not_under_bottom = interpolated_image >= 0
|
|
under_bottom_range = -mn
|
|
under_bottom_scale = under_bottom_target_range / under_bottom_range
|
|
# scale the under the bottom values to the target range
|
|
|
|
#print("Max value:", interpolated_image.max())
|
|
#print("Min value:", interpolated_image.min())
|
|
|
|
interpolated_image[under_bottom] = interpolated_image[under_bottom] * under_bottom_scale
|
|
#self.visualize_tensor(torch.from_numpy(interpolated_image), "After Under-the-Bottom Compression")
|
|
|
|
# scale and shift the non-under-the-bottom to make room
|
|
#interpolated_image[not_under_bottom] = interpolated_image[not_under_bottom] * (
|
|
# 1 - over_top_target_range) + under_bottom_target_range
|
|
|
|
#self.visualize_tensor(torch.from_numpy(interpolated_image),"After Non-Under-the-Bottom Scaling and Shifting")
|
|
|
|
#print("Max value:", interpolated_image.max())
|
|
#print("Min value:", interpolated_image.min())
|
|
|
|
# clamp
|
|
interpolated_image = np.clip(interpolated_image, 0, 1)
|
|
|
|
|
|
else:
|
|
pass
|
|
|
|
# Convert the resized image back to a PyTorch tensor
|
|
# interpolated_image = torch.from_numpy(interpolated_image).to(device, dtype=torch.float16).permute(2, 0, 1)
|
|
interpolated_image = torch.from_numpy(interpolated_image)
|
|
|
|
# convert fix range to 0-255
|
|
interpolated_image = interpolated_image * 255
|
|
# convert to uint8
|
|
interpolated_image = interpolated_image.to(device, dtype=torch.uint8).permute(2, 0, 1)
|
|
|
|
interpolated_images.append(interpolated_image)
|
|
del image
|
|
del image_np
|
|
del interpolated_image
|
|
|
|
# transform this list of tensors into a single tensor of shape (batch_size, channels, height, width)
|
|
|
|
# pre allocate a tensor of shape (batch_size, channels, target_height, target_width)
|
|
y = torch.zeros((x.shape[0], interpolated_images[0].shape[0], self.target_size[1], self.target_size[0]),
|
|
device=device,
|
|
dtype=interpolated_images[0].dtype)
|
|
y = torch.stack(interpolated_images, dim=0, out=y)
|
|
del interpolated_images
|
|
|
|
return y
|
|
|
|
|
|
class LanczosInterpolationModel(nn.Module):
|
|
"""
|
|
A PyTorch model that applies Lanczos interpolation to resize a batch of images.
|
|
|
|
Args:
|
|
target_size (tuple): The target size for resizing the images, in the format (width, height).
|
|
"""
|
|
|
|
def __init__(self, target_size):
|
|
super(LanczosInterpolationModel, self).__init__()
|
|
self.interpolation_layer = LanczosInterpolationLayer(target_size)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Performs the forward pass of the Lanczos interpolation model.
|
|
|
|
Args:
|
|
x (torch.Tensor): A batch of input images with shape (batch_size, channels, height, width).
|
|
|
|
Returns:
|
|
torch.Tensor: A batch of resized images with shape (batch_size, channels, target_height, target_width).
|
|
"""
|
|
return self.interpolation_layer(x)
|
|
|
|
|
|
def expand_mask(mask, expand, tapered_corners, device=None, dtype=None):
|
|
"""
|
|
Convenience function for expanding or eroding a mask using the MaskExpansionModel.
|
|
|
|
Args:
|
|
mask (torch.Tensor): The input mask tensor.
|
|
expand (int): The number of iterations to expand (positive) or erode (negative) the mask.
|
|
tapered_corners (bool): Whether to use tapered corners when expanding or eroding the mask.
|
|
device (torch.device, optional): The device to use for the model and input tensor.
|
|
If not provided, it will use the default device (CPU).
|
|
dtype (torch.dtype, optional): The data type to use for the input tensor.
|
|
If not provided, it will use the default dtype (float32).
|
|
|
|
Returns:
|
|
torch.Tensor: The expanded or eroded mask tensor.
|
|
"""
|
|
# Create an instance of the MaskExpansionModel
|
|
model = MaskExpansionModel()
|
|
|
|
# Move the model to the specified device (or default to CPU)
|
|
device = device or torch.device("cpu")
|
|
model = model.to(device)
|
|
|
|
# Move the input tensor to the specified device and convert to the specified dtype
|
|
dtype = dtype or torch.float32
|
|
mask = mask.to(device, dtype=dtype)
|
|
|
|
# Expand or erode the mask using the MaskExpansionModel
|
|
with torch.no_grad():
|
|
expanded_mask = model(mask, expand, tapered_corners)
|
|
|
|
# Clear the GPU cache if using a GPU
|
|
if device != torch.device("cpu"):
|
|
torch.cuda.empty_cache()
|
|
|
|
return expanded_mask
|
|
|
|
|
|
def lanczos_resize(image, target_size, device="cpu", dtype=torch.float32, output_type="tensor"):
|
|
"""
|
|
Convenience function for resizing an image using the Lanczos interpolation model.
|
|
|
|
Args:
|
|
image (PIL.Image or torch.Tensor): The input image to be resized.
|
|
target_size (tuple): The target size for resizing the image, in the format (width, height).
|
|
device (torch.device, optional): The device to use for the model and input tensor.
|
|
If not provided, it will use the default device (CPU).
|
|
dtype (torch.dtype, optional): The data type to use for the input tensor.
|
|
If not provided, it will use the default dtype (float32).
|
|
output_type (str, optional): The type of the output image.
|
|
If "tensor", the function will return a PyTorch tensor.
|
|
If "pil", the function will return a PIL Image.
|
|
If not provided, it will return a tensor by default.
|
|
|
|
Returns:
|
|
torch.Tensor or PIL.Image: The resized image.
|
|
"""
|
|
import gc
|
|
import os
|
|
import sys
|
|
if sys.platform == "win32":
|
|
import psutil
|
|
print(f"Free memory: {psutil.virtual_memory().available / 1024 / 1024} MB")
|
|
else:
|
|
# works on linux and mac
|
|
print(f"Free memory: {os.popen('free -m').readlines()[1].split()[3]} MB")
|
|
|
|
# Create an instance of the LanczosInterpolationModel
|
|
model = LanczosInterpolationModel(target_size)
|
|
|
|
# Move the model to the specified device (or default to CPU)
|
|
device = device or torch.device("cpu")
|
|
model = model.to(device)
|
|
|
|
# Convert the input image to a PyTorch tensor
|
|
if isinstance(image, list):
|
|
if isinstance(image[0], Image.Image):
|
|
try:
|
|
images = np.array(image)
|
|
except:
|
|
images = np.array([np.array(img) for img in image])
|
|
|
|
images = torch.from_numpy(images).to(device, dtype).permute(0, 3, 1, 2)
|
|
elif isinstance(image, torch.Tensor):
|
|
images = image.to(device, dtype)
|
|
|
|
del image
|
|
torch.cuda.empty_cache()
|
|
|
|
# Resize the image using the Lanczos interpolation model
|
|
with torch.no_grad():
|
|
resized_images = model(images)
|
|
|
|
# Convert the resized tensor back to a PIL Image or keep it as a tensor based on the output_type
|
|
if output_type.lower() == "pil":
|
|
resized_images = resized_images.squeeze(0).permute(0, 2, 3, 1).cpu().numpy()
|
|
resized_images = Image.fromarray(resized_images.astype(np.uint8))
|
|
elif output_type.lower() == "tensor":
|
|
pass # keep the resized_image as a tensor
|
|
else:
|
|
raise ValueError(f"Invalid output_type: {output_type}. Valid options are 'tensor' or 'pil'.")
|
|
|
|
# Clear the GPU cache if using a GPU
|
|
if device != torch.device("cpu"):
|
|
torch.cuda.empty_cache()
|
|
|
|
del images
|
|
del model
|
|
gc.collect()
|
|
# print current free memory on cpu, in a way that is cross platform 1 way that works on all
|
|
|
|
if sys.platform == "win32":
|
|
import psutil
|
|
print(f"Free memory: {psutil.virtual_memory().available / 1024 / 1024} MB")
|
|
else:
|
|
# works on linux and mac
|
|
print(f"Free memory: {os.popen('free -m').readlines()[1].split()[3]} MB")
|
|
|
|
return resized_images
|
|
|
|
|
|
def test_lanczos_interpolation_model():
|
|
"""
|
|
Tests the Lanczos interpolation model by resizing a batch of random images and benchmarking the performance.
|
|
"""
|
|
# Set the target sizes for resizing
|
|
hd_size = (1920, 1080)
|
|
sd_size = (1280, 720)
|
|
|
|
# Generate random PIL images
|
|
batch_size = 600
|
|
hd_images = [Image.fromarray(np.random.randint(0, 256, size=(hd_size[1], hd_size[0], 3), dtype=np.uint8))
|
|
for _ in range(batch_size)]
|
|
sd_images = [Image.fromarray(np.random.randint(0, 256, size=(sd_size[1], sd_size[0], 3), dtype=np.uint8))
|
|
for _ in range(batch_size)]
|
|
|
|
# Perform the benchmarks
|
|
num_iterations = 3
|
|
|
|
# HD to SD benchmark
|
|
start_time = time.time()
|
|
for _ in range(num_iterations):
|
|
sd_output_images = lanczos_resize(hd_images, sd_size)
|
|
end_time = time.time()
|
|
hd_to_sd_elapsed_time = end_time - start_time
|
|
hd_to_sd_avg_time_per_iteration = hd_to_sd_elapsed_time / num_iterations
|
|
|
|
# SD to HD benchmark
|
|
start_time = time.time()
|
|
for _ in range(num_iterations):
|
|
hd_output_images = lanczos_resize(sd_images, hd_size)
|
|
end_time = time.time()
|
|
sd_to_hd_elapsed_time = end_time - start_time
|
|
sd_to_hd_avg_time_per_iteration = sd_to_hd_elapsed_time / num_iterations
|
|
|
|
# Check the sizes of the output images
|
|
|
|
assert sd_output_images.shape[0] == batch_size
|
|
|
|
assert hd_output_images.shape[0] == batch_size
|
|
|
|
print(f"Lanczos interpolation test passed!")
|
|
print(f"Benchmark results:")
|
|
print(f"Number of iterations: {num_iterations}")
|
|
print(f"HD to SD:")
|
|
print(f" Total elapsed time: {hd_to_sd_elapsed_time:.4f} seconds")
|
|
print(f" Average time per iteration: {hd_to_sd_avg_time_per_iteration:.4f} seconds")
|
|
print(f" Average time per HD image -> SD image: {hd_to_sd_avg_time_per_iteration / batch_size:.4f} seconds")
|
|
print(f"SD to HD:")
|
|
print(f" Total elapsed time: {sd_to_hd_elapsed_time:.4f} seconds")
|
|
print(f" Average time per iteration: {sd_to_hd_avg_time_per_iteration:.4f} seconds")
|
|
print(f" Average time per SD image -> HD image: {sd_to_hd_avg_time_per_iteration / batch_size:.4f} seconds")
|
|
|
|
|
|
def test_mask_expansion_model():
|
|
"""
|
|
Tests the MaskExpansionModel by expanding and eroding random masks and benchmarking the performance.
|
|
"""
|
|
# Generate random masks
|
|
batch_size = 1000
|
|
mask_size = (512, 512)
|
|
# masks are 0-1 tensors float
|
|
masks = torch.rand((batch_size, mask_size[0], mask_size[1])).to("cuda")
|
|
print(masks.min(), masks.max())
|
|
|
|
# Perform the benchmarks
|
|
num_iterations = 1
|
|
expand_steps = 5
|
|
erode_steps = -5
|
|
|
|
# Mask expansion benchmark
|
|
start_time = time.time()
|
|
for _ in range(num_iterations):
|
|
print(masks.mean())
|
|
expanded_masks = expand_mask(masks, expand_steps, tapered_corners=True, device="cuda")
|
|
print(masks.mean())
|
|
end_time = time.time()
|
|
expansion_elapsed_time = end_time - start_time
|
|
expansion_avg_time_per_iteration = expansion_elapsed_time / num_iterations
|
|
|
|
# Mask erosion benchmark
|
|
start_time = time.time()
|
|
for _ in range(num_iterations):
|
|
print(masks.mean())
|
|
eroded_masks = expand_mask(masks, erode_steps, tapered_corners=True, device="cuda")
|
|
print(eroded_masks.mean())
|
|
end_time = time.time()
|
|
erosion_elapsed_time = end_time - start_time
|
|
erosion_avg_time_per_iteration = erosion_elapsed_time / num_iterations
|
|
|
|
# Check the sizes of the output masks
|
|
expected_mask_size = (batch_size, mask_size[0], mask_size[1])
|
|
assert expanded_masks.shape == expected_mask_size, f"Expected expanded mask size: {expected_mask_size}, but got: {expanded_masks.shape}"
|
|
assert eroded_masks.shape == expected_mask_size, f"Expected eroded mask size: {expected_mask_size}, but got: {eroded_masks.shape}"
|
|
|
|
print(f"Mask expansion/erosion test passed!")
|
|
print(f"Benchmark results:")
|
|
print(f"Number of iterations: {num_iterations}")
|
|
print(f"Mask expansion:")
|
|
print(f" Total elapsed time: {expansion_elapsed_time:.4f} seconds")
|
|
print(f" Average time per iteration: {expansion_avg_time_per_iteration:.4f} seconds")
|
|
print(f" Average time per mask expansion: {expansion_avg_time_per_iteration / batch_size:.4f} seconds")
|
|
print(f"Mask erosion:")
|
|
print(f" Total elapsed time: {erosion_elapsed_time:.4f} seconds")
|
|
print(f" Average time per iteration: {erosion_avg_time_per_iteration:.4f} seconds")
|
|
print(f" Average time per mask erosion: {erosion_avg_time_per_iteration / batch_size:.4f} seconds")
|
|
|
|
|
|
if __name__ == "image_utils":
|
|
print("image_utils module loaded.")
|
|
print("Testing Lanczos interpolation model...")
|
|
test_lanczos_interpolation_model()
|
|
print("Testing Mask expansion model...")
|
|
# test_mask_expansion_model()
|
|
exit(0)
|
|
|
|
if __name__ == "__main__":
|
|
test_lanczos_interpolation_model()
|
|
# test_mask_expansion_model()
|