FFT/2d_fourie/extractors.py
2023-03-06 11:01:04 +01:00

193 lines
5.8 KiB
Python

import numpy as np
from scipy.spatial import Voronoi
import cv2
class Image_Wrapper:
# def __init__(self, img, x_lower, x_res, y_lower, y_res):
# self.img = img
# self.x_lower = x_lower
# self.y_lower = y_lower
# self.x_res = x_res
# self.y_res = y_res
# self.x_upper = self.x_lower + self.img.shape[0]*self.x_res - self.x_res
# self.y_upper = self.y_lower + self.img.shape[1]*self.y_res - self.x_res
def __init__(self, img, fx, fy):
self.img = img
self.x_lower = np.min(fx)
self.y_lower = np.min(fy)
self.x_upper = np.max(fx)
self.y_upper = np.max(fy)
self.x_res = (self.x_upper - self.x_lower) / self.img.shape[0]
self.y_res = (self.y_upper - self.y_lower) / self.img.shape[0]
def val2pos(self, x, y):
x = (x - self.x_lower) / self.x_res
y = (y - self.y_lower) / self.y_res
return x, y
def check_bounds(self, xl, yl, xu, yu):
if xl > self.img.shape[0]:
print("xl lim")
return False
if yl > self.img.shape[1]:
print("yl lim")
return False
if xu < 0:
print("xu lim")
return False
if yu < 0:
print("yu lim")
return False
return True
def clean_bounds(self, xl, yl, xu, yu):
if xl < 0:
xl = 0
if yl < 0:
yl = 0
if xu > self.img.shape[0]:
xu = self.img.shape[0]
if yu > self.img.shape[1]:
yu = self.img.shape[1]
return xl, yl, xu, yu
def ext(self):
return [self.x_lower, self.x_upper, self.y_lower, self.y_upper]
class Evaluator:
def extract(self, img: Image_Wrapper):
all_val = []
all_idx = []
for ev_points in self.eval_points:
temp_val = []
temp_idx = []
for num in ev_points:
if np.sum(self.mask == num) == 0:
continue
temp_val.append(np.mean(img.img[self.mask == num]))
temp_idx.append(num)
all_val.append(temp_val)
all_idx.append(temp_idx)
all_val.append([np.mean(img.img[self.mask == -1])])
all_idx.append([-1])
return all_idx, all_val
def debug(self, img: Image_Wrapper):
count = 1
for ev_points in self.eval_points:
if count != 1 and False:
count += 1
continue
for num in ev_points:
img.img[self.mask == num] += 10000 * num
count += 1
return img.img
def get_mask(self):
return self.mask
def merge_mask(self):
new_eval_points = np.arange(len(self.eval_points))
new_eval = []
for nc, ev_points in zip(new_eval_points, self.eval_points):
new_eval.append([nc])
for num in ev_points:
self.mask[self.mask == num] = nc
self.eval_points = new_eval
class Voronoi_Evaluator(Evaluator):
def __init__(self, list_points):
points = np.concatenate(list_points, axis=0)
self.eval_points = []
start = 0
for l in list_points:
stop = l.shape[0]
self.eval_points.append(np.arange(start, start + stop))
start += stop
self.vor = Voronoi(points)
def generate_mask(self, img: Image_Wrapper):
self.mask = np.full_like(img.img, -1)
counter = -1
region_mask = self.vor.point_region
for i in np.array(self.vor.regions, dtype=list)[region_mask]:
counter += 1
if -1 in i:
continue
if len(i) == 0:
continue
pts = self.vor.vertices[i]
pts = np.stack(img.val2pos(
pts[:, 0], pts[:, 1])).astype(np.int32).T
if np.any(pts < 0):
continue
mask = np.zeros_like(img.img)
cv2.fillConvexPoly(mask, pts, 1)
mask = mask > 0 # To convert to Boolean
self.mask[mask] = counter
class Rect_Evaluator(Evaluator):
# def __init__(self, points, eval_points):
# self.eval_points = eval_points
# self.points = points
# self.length = 4
def __init__(self, list_points, length=4):
self.points = np.concatenate(list_points, axis=0)
self.eval_points = []
start = 0
for l in list_points:
stop = l.shape[0]
self.eval_points.append(np.arange(start, start + stop))
start += stop
print(start)
print(self.points.shape)
self.length = length
def generate_mask(self, img: Image_Wrapper):
self.mask = np.full_like(img.img, -1)
count = 0
for x, y in self.points:
x, y = img.val2pos(x, y)
x_lower = int(x - self.length)
y_lower = int(y - self.length)
x_upper = int(x + self.length + 1)
y_upper = int(y + self.length + 1)
if img.check_bounds(x_lower, y_lower, x_upper, y_upper):
x_lower, y_lower, x_upper, y_upper = img.clean_bounds(
x_lower, y_lower, x_upper, y_upper)
self.mask[y_lower:y_upper, x_lower:x_upper] = count
count += 1
return all
#
# def main():
# np.random.seed(10)
# points = (np.random.rand(100, 2)-0.5) * 2
# voro = Voronoi_Evaluator(points, [[1],[2]])
# rect = Rect_Evaluator(points, [[1], [2]])
# Z = np.ones((1000, 1000))
# img = Image_Wrapper(Z, -5, .01, -5, .01)
# voro.extract(img)
# rect.extract(img)
#
# plt.scatter(points[[1], 0], points[[1], 1])
# plt.scatter(points[[2], 0], points[[2], 1])
# plt.imshow(img.img, extent=img.ext(), origin="lower")
# #plt.imshow(img.img, origin="lower")
# plt.show()
#
#
# if __name__ == "__main__":
# main()