162 lines
5.1 KiB
Python
162 lines
5.1 KiB
Python
import numpy as np
|
|
import tqdm
|
|
# from scipy.spatial import Voronoi
|
|
# import cv2
|
|
from cache import persist_to_file, timeit
|
|
from spin_image import FFT
|
|
from tools import clean_bounds_offset
|
|
from abc import abstractmethod, ABC
|
|
import logging
|
|
logger = logging.getLogger("fft")
|
|
|
|
|
|
class Evaluator(ABC):
|
|
eval_points: list[np.ndarray]
|
|
|
|
def extract(self, img: FFT):
|
|
all_val = []
|
|
all_idx = []
|
|
for ev_points in self.eval_points:
|
|
logger.info(f"Extracting points: {ev_points}")
|
|
temp_val = []
|
|
temp_idx = []
|
|
for num in ev_points:
|
|
if np.sum(self.mask == num) == 0:
|
|
continue
|
|
temp_val.append(np.sum(img.intens[self.mask == num]))
|
|
temp_idx.append(num)
|
|
all_val.append(temp_val)
|
|
all_idx.append(temp_idx)
|
|
all_val.append([np.sum(img.intens[self.mask == -1])])
|
|
all_idx.append([-1])
|
|
return all_idx, all_val
|
|
|
|
def debug(self, img: FFT):
|
|
for count, ev_points in enumerate(self.eval_points, start=0):
|
|
for num in ev_points:
|
|
img.intens[self.mask == num] += count
|
|
return img
|
|
|
|
def get_mask(self):
|
|
return self.mask
|
|
|
|
@timeit
|
|
def generate_mask(self, img: FFT, merge=False):
|
|
hash = str(img.intens.shape)
|
|
self.mask = self.gen_mask_helper(img)
|
|
if merge:
|
|
self.mask = self.merge_mask_helper()
|
|
self.eval_points = [[a] for a in np.arange(len(self.eval_points))]
|
|
|
|
@abstractmethod
|
|
def gen_mask_helper(self, img: FFT, hash=str):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def merge_mask_helper(self, hash: str):
|
|
pass
|
|
|
|
def purge(self, img: FFT):
|
|
img.intens[self.mask != -1] = 0
|
|
|
|
|
|
class Rect_Evaluator(Evaluator):
|
|
# def __init__(self, points, eval_points):
|
|
# self.eval_points = eval_points
|
|
# self.points = points
|
|
# self.length = 4
|
|
def __init__(self, spots: list[tuple[np.ndarray, np.ndarray]], length: int = 6):
|
|
self.spots = spots
|
|
self.length = length
|
|
self.eval_points = []
|
|
start = 0
|
|
for sp in spots:
|
|
self.eval_points.append(np.arange(start, start+sp[0].size))
|
|
start += sp[0].size
|
|
|
|
def merge_mask_helper(self):
|
|
new_eval_points = np.arange(len(self.eval_points))
|
|
mask = self.mask
|
|
for nc, ev_points in zip(new_eval_points, self.eval_points):
|
|
maske_low = np.min(ev_points) >= self.mask
|
|
maske_high = np.max(ev_points) <= self.mask
|
|
mask[np.logical_and(maske_high, maske_low)] = nc
|
|
return mask
|
|
|
|
def gen_mask_helper(self, img: FFT):
|
|
mask = np.full_like(img.intens, -1)
|
|
count = 0
|
|
for spot_group in self.spots:
|
|
logger.debug(f"Spot: {spot_group}")
|
|
(x, y) = spot_group
|
|
x, y = img.val2pos(x, y)
|
|
for x_p, y_p in zip(x, y):
|
|
xl, yl, xu, yu = clean_bounds_offset(
|
|
x_p, y_p, self.length, img.intens.shape)
|
|
mask[xl:xu, yl:yu] = count
|
|
count += 1
|
|
return mask
|
|
|
|
#
|
|
# 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()
|
|
# 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)
|
|
#
|
|
# @persist_to_file("cache_merge_voro")
|
|
# def merge_mask_helper(self):
|
|
# new_eval_points = np.arange(len(self.eval_points))
|
|
# mask = self.mask
|
|
# for nc, ev_points in zip(new_eval_points, self.eval_points):
|
|
# for num in ev_points:
|
|
# mask[self.mask == num] = nc
|
|
# return mask
|
|
#
|
|
# @persist_to_file("cache_voro")
|
|
# def gen_mask_helper(self, img: Image_Wrapper):
|
|
# 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_2 = np.zeros_like(img.img)
|
|
# cv2.fillConvexPoly(mask_2, pts, 1)
|
|
# mask_2 = mask_2 > 0 # To convert to Boolean
|
|
# mask[mask_2] = counter
|
|
# return mask
|