init.py

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]:
        print("Starting separation...")
        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):
        print("Cleaning planes...")
        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):
        print("Starting normal process...")
        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