| import os |
| from typing import Tuple |
|
|
| import cv2 |
| import numpy.typing as npt |
|
|
| from config import ( |
| NormalMapGenerator, |
| Planes, |
| RGBFlatten, |
| artifactFlatten, |
| grayscaleInv, |
| heightCreation, |
| kernels, |
| stampFill, |
| valueFill, |
| ) |
|
|
|
|
| class Init_Planes: |
| def __init__(self, base_image: npt.NDArray, depth_image: npt.NDArray) -> None: |
| self.base_image = base_image |
| self.depth_image = depth_image |
|
|
| def separation(self, depth: npt.NDArray) -> Tuple[npt.NDArray, npt.NDArray, npt.NDArray]: |
| |
| contrast = cv2.convertScaleAbs(depth, alpha=3, beta=-100) |
| colorIndex = RGBFlatten(contrast).initiate() |
|
|
| mask_list = [] |
| for i in range(3): |
| mask_img = cv2.cvtColor(Planes(colorIndex, i).create_planes(), cv2.COLOR_RGB2BGR) |
| mask_list.append(mask_img) |
|
|
| return mask_list |
|
|
| def data_augmentation(self, mask_list: list): |
| |
| cleaned_segments = [] |
|
|
| for i in range(3): |
| fill_img = stampFill(mask_list[i], i).run() |
| clean_img = valueFill(fill_img).run() |
| cleaned_segments.append(clean_img) |
|
|
| return cleaned_segments |
|
|
| def normal_process(self, img: npt.NDArray): |
| |
| grayscale = grayscaleInv(img).run() |
| heightMap = heightCreation(grayscale, scale=0.5).run() |
| blurred_image_list = artifactFlatten(heightMap).run() |
| normal_maps = [] |
|
|
| for image in blurred_image_list: |
| normal_map = NormalMapGenerator(image).generate_normal_map() |
| normal_maps.append(normal_map) |
| |
| return normal_maps |
|
|
| def initiate_split(self) -> None: |
| mask_list = self.separation(self.depth_image) |
| cleaned_segments = self.data_augmentation(mask_list) |
| normal_maps = self.normal_process(self.base_image) |
| return cleaned_segments, normal_maps |
|
|
