324 lines
12 KiB
Python
324 lines
12 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 = None
|
|
# A list of common font names for cross-platform compatibility
|
|
font_names = [
|
|
"arial.ttf", # Windows, macOS
|
|
"DejaVuSans.ttf", # Common on Linux
|
|
"LiberationSans-Regular.ttf" # Another common Linux font
|
|
]
|
|
|
|
for font_name in font_names:
|
|
try:
|
|
font = ImageFont.truetype(font_name, size)
|
|
break # Stop searching once a font is found
|
|
except IOError:
|
|
continue # Try the next font in the list
|
|
|
|
if font is None:
|
|
# If no common fonts are found, fall back to the default bitmap font.
|
|
print(f"Could not find any of {font_names}. Falling back to default font. Text may be very small.")
|
|
font = ImageFont.load_default()
|
|
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
# Use textbbox for better centering if available (Pillow >= 8.0.0)
|
|
if hasattr(draw, 'textbbox'):
|
|
bbox = draw.textbbox((0, 0), text, font=font)
|
|
text_width = bbox[2] - bbox[0]
|
|
text_height = bbox[3] - bbox[1]
|
|
x = (img.width - text_width) / 2
|
|
y = (img.height - text_height) / 2
|
|
# Adjust position based on the bounding box's top-left corner
|
|
text_position = (x - bbox[0], y - bbox[1])
|
|
else: # Fallback for older Pillow versions
|
|
# textsize is deprecated but works as a fallback.
|
|
text_width, text_height = draw.textsize(text, font=font)
|
|
x = (img.width - text_width) / 2
|
|
y = (img.height - text_height) / 2
|
|
text_position = (x, y)
|
|
|
|
# Draw outline for better visibility
|
|
outline_color = (0, 0, 0)
|
|
for offset in [(-1, -1), (-1, 1), (1, -1), (1, 1), (-1, 0), (1, 0), (0, -1), (0, 1)]:
|
|
draw.text((text_position[0] + offset[0], text_position[1] + offset[1]), text, font=font, fill=outline_color)
|
|
|
|
# Draw the main 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
|
|
|
|
# --- Image Difference Calculation Functions ---
|
|
|
|
def calculate_mse_diff(frames1_t, frames2_t):
|
|
"""Calculates the Mean Squared Error between two batches of frames."""
|
|
diffs = frames1_t.float() - frames2_t.float()
|
|
return torch.mean(diffs ** 2, dim=(1, 2, 3))
|
|
|
|
def calculate_ssim_diff(frames1_p, frames2_p):
|
|
"""Calculates the Structural Similarity Index (SSIM) difference between two lists of PIL images."""
|
|
from skimage.metrics import structural_similarity as ssim
|
|
diffs = []
|
|
for i in range(len(frames1_p)):
|
|
# Convert images to grayscale numpy arrays for SSIM
|
|
im1 = np.array(frames1_p[i].convert('L'))
|
|
im2 = np.array(frames2_p[i].convert('L'))
|
|
# SSIM returns a value between -1 and 1, where 1 is perfect similarity.
|
|
# We want the difference, so we use 1 - ssim.
|
|
similarity = ssim(im1, im2, data_range=im2.max() - im2.min())
|
|
diffs.append(1.0 - similarity)
|
|
return torch.tensor(diffs, device=device)
|
|
|
|
def calculate_phash_diff(frames1_p, frames2_p):
|
|
"""Calculates the perceptual hash (pHash) difference between two lists of PIL images."""
|
|
import imagehash
|
|
diffs = []
|
|
for i in range(len(frames1_p)):
|
|
hash1 = imagehash.phash(frames1_p[i])
|
|
hash2 = imagehash.phash(frames2_p[i])
|
|
# The difference is the number of bits that are different (Hamming distance).
|
|
# We normalize it by the total number of bits in the hash (usually 64).
|
|
diffs.append((hash1 - hash2) / len(hash1))
|
|
return torch.tensor(diffs, device=device)
|
|
|
|
def calculate_contrast(frame_p):
|
|
"""Calculates the RMS contrast of a single PIL image."""
|
|
# Convert to grayscale numpy array
|
|
im = np.array(frame_p.convert('L'), dtype=float)
|
|
return im.std()
|
|
|
|
def calculate_edge_diff(frame1_p, frame2_p, focus_ratio=0.3):
|
|
"""Calculates the difference between the edges of the central focused area of two PIL images."""
|
|
import cv2
|
|
|
|
def process_frame(frame):
|
|
# Crop to central focus area
|
|
w, h = frame.size
|
|
crop_w, crop_h = int(w * focus_ratio), int(h * focus_ratio)
|
|
left, top = (w - crop_w) // 2, (h - crop_h) // 2
|
|
right, bottom = left + crop_w, top + crop_h
|
|
cropped = frame.crop((left, top, right, bottom))
|
|
|
|
# Convert to grayscale numpy array
|
|
gray = np.array(cropped.convert('L'))
|
|
|
|
# Apply Laplacian operator to detect edges
|
|
edges = cv2.Laplacian(gray, cv2.CV_64F)
|
|
return edges
|
|
|
|
edges1 = process_frame(frame1_p)
|
|
edges2 = process_frame(frame2_p)
|
|
|
|
# Calculate Mean Squared Error between the edge maps
|
|
return np.mean((edges1 - edges2) ** 2)
|