File size: 4,231 Bytes
dd06d6b |
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 |
try:
import cv2
except:
pass
import numpy as np
from PIL import Image
from relighting.ball_processor import get_ideal_normal_ball
def create_grid(image_size, n_ball, size):
height, width = image_size
nx, ny = n_ball
if nx * ny == 1:
grid = np.array([[(height-size)//2, (width-size)//2]])
else:
height_ = np.linspace(0, height-size, nx).astype(int)
weight_ = np.linspace(0, width-size, ny).astype(int)
hh, ww = np.meshgrid(height_, weight_)
grid = np.stack([hh,ww], axis = -1).reshape(-1,2)
return grid
class MaskGenerator():
def __init__(self, cache_mask=True):
self.cache_mask = cache_mask
self.all_masks = []
def clear_cache(self):
self.all_masks = []
def retrieve_masks(self):
return self.all_masks
def generate_grid(self, image, mask_ball, n_ball=16, size=128):
ball_positions = create_grid(image.size, n_ball, size)
# _, mask_ball = get_normal_ball(size)
masks = []
mask_template = np.zeros(image.size)
for x, y in ball_positions:
mask = mask_template.copy()
mask[y:y+size, x:x+size] = 255 * mask_ball
mask = Image.fromarray(mask.astype(np.uint8), "L")
masks.append(mask)
# if self.cache_mask:
# self.all_masks.append((x, y, size))
return masks, ball_positions
def generate_single(self, image, mask_ball, x, y, size):
w,h = image.size # numpy as (h,w) but PIL is (w,h)
mask = np.zeros((h,w))
mask[y:y+size, x:x+size] = 255 * mask_ball
mask = Image.fromarray(mask.astype(np.uint8), "L")
return mask
def generate_best(self, image, mask_ball, size):
w,h = image.size # numpy as (h,w) but PIL is (w,h)
mask = np.zeros((h,w))
(y, x), _ = find_best_location(np.array(image), ball_size=size)
mask[y:y+size, x:x+size] = 255 * mask_ball
mask = Image.fromarray(mask.astype(np.uint8), "L")
return mask, (x, y)
def get_only_high_freqency(image: np.array):
"""
Get only height freqency image by subtract low freqency (using gaussian blur)
@params image: np.array - image in RGB format [h,w,3]
@return high_frequency: np.array - high freqnecy image in grayscale format [h,w]
"""
# Convert to grayscale
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
# Subtract low freqency from high freqency
kernel_size = 11 # Adjust this according to your image size
high_frequency = gray - cv2.GaussianBlur(gray,(kernel_size, kernel_size), 0)
return high_frequency
def find_best_location(image, ball_size=128):
"""
Find the best location to place the ball (Eg. empty location)
@params image: np.array - image in RGB format [h,w,3]
@return min_pos: tuple - top left position of the best location (the location is in "Y,X" format)
@return min_val: float - the sum value contain in the window
"""
local_variance = get_only_high_freqency(image)
qsum = quicksum2d(local_variance)
min_val = None
min_pos = None
k = ball_size
for i in range(k-1, qsum.shape[0]):
for j in range(k-1, qsum.shape[1]):
A = 0 if i-k < 0 else qsum[i-k, j]
B = 0 if j-k < 0 else qsum[i, j-k]
C = 0 if (i-k < 0) or (j-k < 0) else qsum[i-k, j-k]
sum = qsum[i, j] - A - B + C
if (min_val is None) or (sum < min_val):
min_val = sum
min_pos = (i-k+1, j-k+1) # get top left position
return min_pos, min_val
def quicksum2d(x: np.array):
"""
Quick sum algorithm to find the window that have smallest sum with O(n^2) complexity
@params x: np.array - image in grayscale [h,w]
@return q: np.array - quick sum of the image for future seach in find_best_location [h,w]
"""
qsum = np.zeros(x.shape)
for i in range(x.shape[0]):
for j in range(x.shape[1]):
A = 0 if i-1 < 0 else qsum[i-1, j]
B = 0 if j-1 < 0 else qsum[i, j-1]
C = 0 if (i-1 < 0) or (j-1 < 0) else qsum[i-1, j-1]
qsum[i, j] = A + B - C + x[i, j]
return qsum |