File size: 6,369 Bytes
032e687 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
from PIL import Image
import numpy as np
import cv2
def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
best_ratio_diff = float('inf')
best_ratio = (1, 1)
area = width * height
for ratio in target_ratios:
target_aspect_ratio = ratio[0] / ratio[1]
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
if ratio_diff < best_ratio_diff:
best_ratio_diff = ratio_diff
best_ratio = ratio
elif ratio_diff == best_ratio_diff:
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
best_ratio = ratio
# print(f'width: {width}, height: {height}, best_ratio: {best_ratio}')
return best_ratio
def dynamic_preprocess(image, regions, merged_regions, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
assert image.size == merged_regions.size
orig_width, orig_height = image.size
aspect_ratio = orig_width / orig_height
# calculate the existing image aspect ratio
target_ratios = set(
(i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
i * j <= max_num and i * j >= min_num)
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
# find the closest aspect ratio to the target
target_aspect_ratio = find_closest_aspect_ratio(
aspect_ratio, target_ratios, orig_width, orig_height, image_size)
# calculate the target width and height
target_width = image_size * target_aspect_ratio[0]
target_height = image_size * target_aspect_ratio[1]
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
# resize the image
resized_img = image.resize((target_width, target_height))
resized_merged_regions = merged_regions.resize((target_width, target_height))
# resize the regions
resized_regions = cv2.resize(np.transpose(regions, (1, 2, 0)), dsize=(target_width, target_height), interpolation=cv2.INTER_NEAREST_EXACT)
if resized_regions.ndim < 3:
resized_regions = resized_regions[:, :, np.newaxis]
# for r in range(resized_regions.shape[-1]):
# mask = resized_regions[:, :, r]
# new_img = np.zeros((resized_regions.shape[0], resized_regions.shape[1], 3), dtype=np.uint8)
# new_img[:, :, 0] = mask * 255
# cv2.imwrite(f"./{r}.png", new_img)
processed_images = []
processed_merged_regions = []
processed_regions = [[] for _ in range(resized_regions.shape[-1])]
for i in range(blocks):
box = (
(i % (target_width // image_size)) * image_size,
(i // (target_width // image_size)) * image_size,
((i % (target_width // image_size)) + 1) * image_size,
((i // (target_width // image_size)) + 1) * image_size
)
# split the image
split_img = resized_img.crop(box)
processed_images.append(split_img)
# split the visual prompt canvas
split_mrgn = resized_merged_regions.crop(box)
processed_merged_regions.append(split_mrgn)
split_rgn = resized_regions[box[1]:box[3], box[0]:box[2], :]
for r in range(resized_regions.shape[-1]):
processed_regions[r].append(split_rgn[:, :, r])
assert len(processed_images) == blocks
if use_thumbnail and len(processed_images) != 1:
thumbnail_img = image.resize((image_size, image_size))
processed_images.append(thumbnail_img)
thumbnail_mrng = merged_regions.resize((image_size, image_size))
processed_merged_regions.append(thumbnail_mrng)
thumbnail_rng = cv2.resize(np.transpose(regions, (1, 2, 0)), dsize=(image_size, image_size), interpolation=cv2.INTER_NEAREST_EXACT)
if thumbnail_rng.ndim < 3:
thumbnail_rng = thumbnail_rng[:, :, np.newaxis]
for r in range(regions.shape[0]):
processed_regions[r].append(thumbnail_rng[:, :, r])
# unique_list = []
# for region_i in processed_regions:
# unique_list.append(np.unique(region_i))
return processed_images, processed_regions, processed_merged_regions
# def dynamic_preprocess(image, merged_regions, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
# assert image.size == merged_regions.size
# orig_width, orig_height = image.size
# aspect_ratio = orig_width / orig_height
# # calculate the existing image aspect ratio
# target_ratios = set(
# (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
# i * j <= max_num and i * j >= min_num)
# target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
# # find the closest aspect ratio to the target
# target_aspect_ratio = find_closest_aspect_ratio(
# aspect_ratio, target_ratios, orig_width, orig_height, image_size)
# # calculate the target width and height
# target_width = image_size * target_aspect_ratio[0]
# target_height = image_size * target_aspect_ratio[1]
# blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
# # resize the image
# resized_img = image.resize((target_width, target_height))
# resized_merged_regions = merged_regions.resize((target_width, target_height))
# processed_images = []
# processed_merged_regions = []
# for i in range(blocks):
# box = (
# (i % (target_width // image_size)) * image_size,
# (i // (target_width // image_size)) * image_size,
# ((i % (target_width // image_size)) + 1) * image_size,
# ((i // (target_width // image_size)) + 1) * image_size
# )
# # split the image
# split_img = resized_img.crop(box)
# processed_images.append(split_img)
# # split the visual prompt canvas
# split_mrgn = resized_merged_regions.crop(box)
# processed_merged_regions.append(split_mrgn)
# assert len(processed_images) == blocks
# if use_thumbnail and len(processed_images) != 1:
# thumbnail_img = image.resize((image_size, image_size))
# processed_images.append(thumbnail_img)
# thumbnail_mrng = merged_regions.resize((image_size, image_size))
# processed_merged_regions.append(thumbnail_mrng)
# return processed_images, processed_merged_regions |