Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,199 Bytes
4583098 |
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 |
"""
Helper functions mainly for data preparation.
"""
import cv2
import numpy as np
from PIL import Image
# Extracting the image and mask from gr.ImageEditor.
def process_gradio_source(editor_data):
original_image = editor_data["background"].convert("RGB")
layers = editor_data.get("layers", [])
full_mask = Image.new("L", original_image.size, 0)
cropped_region = original_image.copy()
if not layers:
return full_mask, cropped_region, original_image
try:
layer_image = layers[0]
layer_pos = (0, 0)
if layer_image.mode != "RGBA":
layer_image = layer_image.convert("RGBA")
alpha_channel = layer_image.split()[-1]
full_mask = Image.new("L", original_image.size, 0)
full_mask.paste(alpha_channel, layer_pos)
full_mask = full_mask.point(lambda p: 255 if p > 128 else 0)
original_np = np.array(original_image)
mask_np = np.array(full_mask)
mask_bool = mask_np > 0
cropped_array = np.zeros_like(original_np)
cropped_array[mask_bool] = original_np[mask_bool]
cropped_region = Image.fromarray(cropped_array)
except Exception as e:
print(f"Raise error: {str(e)}")
return original_image, full_mask, cropped_region
# Get bounding box from the mask.
def get_bbox_from_mask(mask_image):
mask_array = np.array(mask_image)
if mask_array.ndim == 3:
mask_array = cv2.cvtColor(mask_array, cv2.COLOR_RGB2GRAY)
_, binary_mask = cv2.threshold(mask_array, 127, 255, cv2.THRESH_BINARY)
contours, _ = cv2.findContours(mask_array, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return None
largest_contour = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(largest_contour)
return (x, y), (x + w, y + h)
# Crop image from the bounding box coord.
def crop_image_from_bb(original_image, top_left, bottom_right):
x1, y1 = map(int, top_left)
x2, y2 = map(int, bottom_right)
if not (0 <= x1 < x2 <= original_image.width and 0 <= y1 < y2 <= original_image.height):
raise ValueError("Invalid bounding box coordinates")
crop_box = (x1, y1, x2, y2)
cropped_image = original_image.crop(crop_box)
return cropped_image
# Resize image to the target size with zero paddings.
def resize_img_and_pad(input_image, target_size):
cropped_width, cropped_height = input_image.size
target_width, target_height = target_size
scale = min(target_width / cropped_width, target_height / cropped_height)
new_width = int(cropped_width * scale)
new_height = int(cropped_height * scale)
resized_image = input_image.resize((new_width, new_height), Image.BILINEAR)
padded_image = Image.new("RGB", target_size, (0, 0, 0))
left_padding = (target_width - new_width) // 2
top_padding = (target_height - new_height) // 2
padded_image.paste(resized_image, (left_padding, top_padding))
return padded_image
def pil_to_np(pil_img):
if pil_img.mode == 'RGBA':
rgb = pil_img.convert('RGB')
alpha = pil_img.split()[3]
img = np.array(rgb)
alpha = np.array(alpha)
else:
img = np.array(pil_img.convert('RGB'))
return img
# Helper function considering the VAE compression factor.
def nearest_multiple_of_16(ref_height):
lower = (ref_height // 16) * 16
upper = ((ref_height + 15) // 16) * 16
if ref_height - lower <= upper - ref_height:
return lower
else:
return upper
# Resize to generate the reference context.
def generate_context_reference_image(reference_image_pil, img_width=512):
reference_image_rgb = pil_to_np(reference_image_pil)
ref_height, ref_width = reference_image_rgb.shape[0], reference_image_rgb.shape[1]
ref_height = int((img_width / ref_width) * ref_height)
ref_height = nearest_multiple_of_16(ref_height)
reference_context = cv2.resize(reference_image_rgb, (img_width, ref_height))
reference_new_pil = Image.fromarray(reference_context)
return reference_new_pil
|