228 lines
7.4 KiB
Python
228 lines
7.4 KiB
Python
import torch
|
|
import numpy as np
|
|
from pathlib import Path
|
|
import os
|
|
from torchvision.transforms import ToTensor, ToPILImage
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
#Get num closest to 8
|
|
def cl8(num):
|
|
rem = num % 8
|
|
if rem <= 4:
|
|
return round(num - rem)
|
|
else:
|
|
return round(num + (8 - rem))
|
|
|
|
def closest_lcm(n, div1, div2):
|
|
# Find the LCM
|
|
lcm = np.lcm(int(div1), int(div2))
|
|
|
|
# Find the multiple of LCM closest to n
|
|
lower_multiple = (n // lcm) * lcm # Largest multiple of lcm less than or equal to n
|
|
upper_multiple = lower_multiple + lcm # Smallest multiple of lcm greater than n
|
|
|
|
# Find which multiple is closer to n
|
|
if n - lower_multiple > upper_multiple - n:
|
|
return upper_multiple
|
|
else:
|
|
return lower_multiple
|
|
|
|
def normalize_size(images):
|
|
refimage = images[0]
|
|
refimage = refimage.resize((cl8(refimage.width), cl8(refimage.height)), Image.Resampling.LANCZOS)
|
|
return_images = []
|
|
for i in range(len(images)):
|
|
if images[i].size != refimage.size:
|
|
images[i] = images[i].resize(refimage.size, Image.Resampling.LANCZOS)
|
|
return_images.append(images[i])
|
|
return return_images
|
|
|
|
def constrain_image(image, max_width, max_height):
|
|
width, height = image.size
|
|
aspect_ratio = width / float(height)
|
|
|
|
if width > max_width or height > max_height:
|
|
if width / float(max_width) > height / float(max_height):
|
|
new_width = max_width
|
|
new_height = int(new_width / aspect_ratio)
|
|
else:
|
|
new_height = max_height
|
|
new_width = int(new_height * aspect_ratio)
|
|
image = image.resize((cl8(new_width), cl8(new_height)), Image.Resampling.LANCZOS)
|
|
|
|
return image
|
|
|
|
def padlist(lst, targetsize):
|
|
if targetsize <= len(lst):
|
|
return lst[:targetsize]
|
|
|
|
last_elem = lst[-1]
|
|
num_repeats = targetsize - len(lst)
|
|
|
|
return lst + [last_elem] * num_repeats
|
|
|
|
def MakeGrid(images, rows, cols):
|
|
widths, heights = zip(*(i.size for i in images))
|
|
|
|
grid_width = max(widths) * cols
|
|
grid_height = max(heights) * rows
|
|
cell_width = grid_width // cols
|
|
cell_height = grid_height // rows
|
|
final_image = Image.new('RGB', (grid_width, grid_height))
|
|
x_offset = 0
|
|
y_offset = 0
|
|
for i in range(len(images)):
|
|
final_image.paste(images[i], (x_offset, y_offset))
|
|
x_offset += cell_width
|
|
if x_offset == grid_width:
|
|
x_offset = 0
|
|
y_offset += cell_height
|
|
|
|
# Save the final image
|
|
return final_image
|
|
|
|
def BreakGrid(grid, rows, cols):
|
|
width = grid.width // cols
|
|
height = grid.height // rows
|
|
outimages = []
|
|
for row in range(rows):
|
|
for col in range(cols):
|
|
left = col * width
|
|
top = row * height
|
|
right = left + width
|
|
bottom = top + height
|
|
current_img = grid.crop((left, top, right, bottom))
|
|
outimages.append(current_img)
|
|
return outimages
|
|
|
|
def ImgLabeler(img, text, size=72, color=(255,255,255)):
|
|
font = ImageFont.truetype("arial.ttf", size)
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
# Get text size
|
|
text_size = draw.textlength(text, font=font)
|
|
# Calculate x, y coordinates of the text
|
|
x = (img.width - text_size) / 2
|
|
y = (img.height - text_size) / 2
|
|
|
|
# Position for the text, centered
|
|
text_position = (x, y)
|
|
|
|
# Draw the text onto the image
|
|
draw.text(text_position, text, font=font, fill=color)
|
|
|
|
return img
|
|
|
|
def load_and_preprocess(image_path):
|
|
with Image.open(image_path) as img:
|
|
return torch.tensor(np.array(img.convert('L')), device=device, dtype=torch.float32)
|
|
|
|
def compute_histogram(tensor):
|
|
# Min and max values for grayscale images
|
|
min_val, max_val = 0, 255
|
|
# Compute the histogram by counting the number of occurrences within each bin
|
|
hist = torch.histc(tensor, bins=255, min=min_val, max=max_val)
|
|
return hist / tensor.numel() # Normalize by the number of elements
|
|
|
|
def get_iterated_path(directory, base_filename, extension='.png'):
|
|
# Construct the full file path
|
|
count = 0
|
|
while True:
|
|
# Append a count to the filename if it's not the first file
|
|
if count == 0:
|
|
unique_filename = f"{base_filename}{extension}"
|
|
else:
|
|
unique_filename = f"{base_filename}_{count}{extension}"
|
|
full_file_path = os.path.join(directory, unique_filename)
|
|
|
|
# Check if a file with this name already exists
|
|
if not os.path.exists(full_file_path):
|
|
break # Exit the loop once the image is saved
|
|
count += 1
|
|
|
|
return full_file_path
|
|
|
|
def CheckMakeDir(dir):
|
|
try:
|
|
if not Path.exists(Path(dir)):
|
|
Path.mkdir(Path(dir), parents=True, exist_ok=True)
|
|
return True, "Success"
|
|
except:
|
|
st_out = f"Error: could not locate nor create directory {dir}."
|
|
print(st_out)
|
|
return False, st_out
|
|
|
|
def apply_conditional_ema_pytorch(images: list[Image.Image], alpha: float, threshold: float) -> list[Image.Image]:
|
|
|
|
processed_images = []
|
|
ema_tensor = None
|
|
to_tensor = ToTensor()
|
|
to_pil = ToPILImage()
|
|
|
|
for image in images:
|
|
frame_tensor = to_tensor(image).to(device)
|
|
|
|
if ema_tensor is None:
|
|
# For the first frame, EMA is just the frame itself
|
|
ema_tensor = frame_tensor.clone()
|
|
else:
|
|
# Calculate EMA for each subsequent frame
|
|
deviation = torch.abs(frame_tensor - ema_tensor)
|
|
update_mask = deviation > threshold # Mask to identify where to update
|
|
|
|
# EMA update
|
|
ema_tensor = alpha * frame_tensor + (1 - alpha) * ema_tensor
|
|
ema_tensor[update_mask] = frame_tensor[update_mask] # Apply conditional update
|
|
|
|
processed_images.append(to_pil(ema_tensor))
|
|
|
|
return processed_images
|
|
|
|
def cat_to_pils(tensor):
|
|
to_pil = ToPILImage()
|
|
tensor = tensor.permute(0, 3, 1, 2)
|
|
pils = [to_pil(tensor[i]) for i in range(tensor.shape[0])]
|
|
return pils
|
|
|
|
def pil_to_tens(pimg):
|
|
to_tensor = ToTensor()
|
|
tensor = to_tensor(pimg).unsqueeze(0).permute(0, 2, 3, 1)
|
|
return tensor
|
|
|
|
def get_grid_aspect(num_images: int, image_width: int, image_height: int) -> (int, int):
|
|
|
|
if num_images == 0:
|
|
return 0, 0
|
|
|
|
min_diff = float('inf')
|
|
best_layout = (1, num_images)
|
|
|
|
if image_width > image_height:
|
|
for cols in range(1, num_images + 1):
|
|
rows = -(-num_images // cols)
|
|
grid_width = cols * image_width
|
|
grid_height = rows * image_height
|
|
diff = abs(grid_width - grid_height)
|
|
|
|
if diff < min_diff:
|
|
min_diff = diff
|
|
best_layout = (rows, cols)
|
|
|
|
if cols > num_images / cols:
|
|
break
|
|
else:
|
|
for rows in range(1, num_images + 1):
|
|
cols = -(-num_images // rows)
|
|
grid_width = cols * image_width
|
|
grid_height = rows * image_height
|
|
diff = abs(grid_height - grid_width)
|
|
|
|
if diff < min_diff:
|
|
min_diff = diff
|
|
best_layout = (rows, cols)
|
|
|
|
if rows > num_images / rows:
|
|
break
|
|
|
|
return best_layout |