csd-solver/color_blindness_processor.py at master · starrdye/csd-solver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import numpy as np
from PIL import Image
import cv2
from typing import Optional
from utils import Transforms, Utils

class ColorBlindnessProcessor:

    @staticmethod
    def simulate(input_path: str,
                 simulate_type: str = 'protanopia',
                 simulate_degree_primary: float = 1.0,
                 simulate_degree_sec: float = 1.0,
                 return_type: str = 'save',
                 save_path: Optional[str] = None):
        """
        Simulate color blindness on an image.

        :param input_path: Input path of the image.
        :param simulate_type: Type of simulation needed. Can be 'protanopia', 'deutranopia', 'tritanopia', 'hybrid'.
        :param simulate_degree_primary: Primary degree of simulation: used for 'protanopia', 'deutranopia', 'tritanopia'.
        :param simulate_degree_sec: Secondary degree of simulation: used for 'hybrid'.
        :param return_type: How to return the Simulated Image. Use 'pil' for PIL.Image, 'np' for Numpy array,
                            'save' for Saving to path.
        :param save_path: Where to save the simulated file. Valid only if return_type is set as 'save'.
        :return: Processed image in the specified format.
        """
        assert simulate_type in ['protanopia', 'deutranopia', 'tritanopia', 'hybrid'], \
            f'Invalid Simulate Type: {simulate_type}'

        # Load the image file in LMS colorspace
        img_lms = Utils.load_lms(input_path)

        if simulate_type == 'protanopia':
            transform = Transforms.lms_protanopia_sim(degree=simulate_degree_primary)
        elif simulate_type == 'deutranopia':
            transform = Transforms.lms_deutranopia_sim(degree=simulate_degree_primary)
        elif simulate_type == 'tritanopia':
            transform = Transforms.lms_tritanopia_sim(degree=simulate_degree_primary)
        else:
            transform = Transforms.hybrid_protanomaly_deuteranomaly_sim(degree_p=simulate_degree_primary,
                                                                        degree_d=simulate_degree_sec)

        # Transforming the LMS Image
        img_sim = np.dot(img_lms, transform)

        # Converting back to RGB colorspace
        img_sim = (np.dot(img_sim, Transforms.lms_to_rgb()) * 255).astype(np.uint8)

        if return_type == 'save':
            assert save_path is not None, 'No save path provided.'
            img_sim_bgr = cv2.cvtColor(img_sim, cv2.COLOR_RGB2BGR)
            cv2.imwrite(save_path, img_sim_bgr)
            return

        if return_type == 'np':
            return img_sim

        if return_type == 'pil':
            return Image.fromarray(img_sim)

    @staticmethod
    def correct(input_path: str,
                protanopia_degree: float = 1.0,
                deutranopia_degree: float = 1.0,
                return_type: str = 'save',
                save_path: Optional[str] = None):
        """
        Correct images for color blindness.

        :param input_path: Input path of the image.
        :param protanopia_degree: Protanopia degree as diagnosed by a doctor using Ishihara test.
        :param deutranopia_degree: Deutranopia degree as diagnosed by a doctor using Ishihara test.
        :param return_type: How to return the corrected image. Use 'pil' for PIL.Image, 'np' for Numpy array,
                            'save' for Saving to path.
        :param save_path: Where to save the corrected file. Valid only if return_type is set as 'save'.
        :return: Processed image in the specified format.
        """
        img_rgb = Utils.load_rgb(input_path)

        transform = Transforms.correction_matrix(protanopia_degree=protanopia_degree,
                                                 deutranopia_degree=deutranopia_degree)

        img_corrected = (np.dot(img_rgb, transform) * 255).astype(np.uint8)

        if return_type == 'save':
            assert save_path is not None, 'No save path provided.'
            img_corrected_bgr = cv2.cvtColor(img_corrected, cv2.COLOR_RGB2BGR)
            cv2.imwrite(save_path, img_corrected_bgr)
            return

        if return_type == 'np':
            return img_corrected

        if return_type == 'pil':
            return Image.fromarray(img_corrected)