| r""" CHM 4D kernel (psi, iso, and full) generator """ | |
| import torch | |
| from .geometry import Geometry | |
| class KernelGenerator: | |
| def __init__(self, ksz, ktype): | |
| self.ksz = ksz | |
| self.idx4d = Geometry.init_idx4d(ksz) | |
| self.kernel = torch.zeros((ksz, ksz, ksz, ksz)) | |
| self.center = (ksz // 2, ksz // 2) | |
| self.ktype = ktype | |
| def quadrant(self, crd): | |
| if crd[0] < self.center[0]: | |
| horz_quad = -1 | |
| elif crd[0] < self.center[0]: | |
| horz_quad = 1 | |
| else: | |
| horz_quad = 0 | |
| if crd[1] < self.center[1]: | |
| vert_quad = -1 | |
| elif crd[1] < self.center[1]: | |
| vert_quad = 1 | |
| else: | |
| vert_quad = 0 | |
| return horz_quad, vert_quad | |
| def generate(self): | |
| return None if self.ktype == 'full' else self.generate_chm_kernel() | |
| def generate_chm_kernel(self): | |
| param_dict = {} | |
| for idx in self.idx4d: | |
| src_i, src_j, trg_i, trg_j = idx | |
| d_tail = Geometry.get_distance((src_i, src_j), self.center) | |
| d_head = Geometry.get_distance((trg_i, trg_j), self.center) | |
| d_off = Geometry.get_distance((src_i, src_j), (trg_i, trg_j)) | |
| horz_quad, vert_quad = self.quadrant((src_j, src_i)) | |
| src_crd = (src_i, src_j) | |
| trg_crd = (trg_i, trg_j) | |
| key = self.build_key(horz_quad, vert_quad, d_head, d_tail, src_crd, trg_crd, d_off) | |
| coord1d = Geometry.get_coord1d((src_i, src_j, trg_i, trg_j), self.ksz) | |
| if param_dict.get(key) is None: param_dict[key] = [] | |
| param_dict[key].append(coord1d) | |
| return param_dict | |
| def build_key(self, horz_quad, vert_quad, d_head, d_tail, src_crd, trg_crd, d_off): | |
| if self.ktype == 'iso': | |
| return '%d' % d_off | |
| elif self.ktype == 'psi': | |
| d_max = max(d_head, d_tail) | |
| d_min = min(d_head, d_tail) | |
| return '%d_%d_%d' % (d_max, d_min, d_off) | |
| else: | |
| raise Exception('not implemented.') | |