utils.py

import cv2
import numpy as np
import tensorflow as tf


class ImageEnhancements:
    def image_enhance_np(self, img: np.ndarray) -> np.ndarray:
        """
        Denoises and applies filter to sharpen edges in numpy images.
        """
        # Clip values for opencv to process
        img_uint8 = np.clip(img * 255, 0, 255).astype(np.uint8)

        # Denoise + sharpen
        deblurred_img = cv2.fastNlMeansDenoisingColored(img_uint8, None, 5, 5, 5, 7)
        sharpen_kernel = np.array(
            [[0, -1, 0], [-1, 6, -1], [0, -1, 0]], dtype=np.float32
        )
        sharpen_img = cv2.filter2D(deblurred_img, -1, sharpen_kernel)

        # Convert back to original data type and pixel range
        enhanced_img = sharpen_img.astype(np.float32) / 255
        enhanced_img = np.clip(enhanced_img, 0, 1).astype(np.float32)

        return enhanced_img

    def image_enhance_tensor(self, img: tf.Tensor) -> tf.Tensor:
        """
        Uses tensorflow wrapper to apply opencv numpy functions to tensors.
        """
        enhanced = tf.numpy_function(
            func=self.image_enhance_np,
            inp=[img],
            Tout=img.dtype,
        )

        # Add shape back to metadata since tf wrapper loses info
        enhanced.set_shape(img.shape)

        return enhanced