tan_specular_highlights.py 30.5 KB
Newer Older
1 2 3 4 5 6
'''
Created on May 9, 2016

@author: sbhatta

#------------------------------#
Amir Mohammadi's avatar
Amir Mohammadi committed
7
   reference:
8 9 10 11
   "separating reflection components of textured surfaces using a single image"
   by Robby T. Tan, Katsushi Ikeuchi,
   IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI),
   27(2), pp.179-193, February, 2005
Amir Mohammadi's avatar
Amir Mohammadi committed
12 13 14

   This Python implementation is inspired by the C++ code provided by Prof.
   Robby Tan:
15 16
       http://tanrobby.github.io/code.html#
       http://tanrobby.github.io/code/highlight.zip
Amir Mohammadi's avatar
Amir Mohammadi committed
17 18 19 20 21

   The main difference between this implementation and the original C++
   implementation is that here the init_labels() function also ignores pixels
   marked G_DIFFUSE. This leads to a smaller number of iterations per epsilon
   value, while producing the same result as the C++ code.
22 23 24 25
#------------------------------#
'''

import numpy as np
Amir Mohammadi's avatar
Amir Mohammadi committed
26
# np.seterr(divide='ignore', invalid='ignore')
27 28 29 30

import bob.io.image
import bob.io.base

Amir Mohammadi's avatar
Amir Mohammadi committed
31 32 33 34 35 36 37 38
# special-pixel-markers. The values don't matter, but they should all
# remain different from each other.
G_SPECULARX = 10
G_SPECULARY = 11
G_DIFFUSE = 12
G_BOUNDARY = 13
G_NOISE = 14
G_CAMERA_DARK = 15
39 40 41


'''
Amir Mohammadi's avatar
Amir Mohammadi committed
42

43
'''
Amir Mohammadi's avatar
Amir Mohammadi committed
44 45


46
def remove_highlights(srcImage, startEps=0.5, verboseFlag=True):
Amir Mohammadi's avatar
Amir Mohammadi committed
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61
    """Returns the separate reflection components (highlights and diffuse
    components) comprising the input color RGB image. This is the main function
    that processes the input srcImage and returns the separate components.

    Args:
        srcImage: numpy array of shape (3, maxY, maxX), containing a RGB image.
            (maxY, maxX) give the image-size (num. rows, num. cols).
        startEps: floating point (small value), determines the number of
                  iterations of algorithm. epsilon is initialized to this
                  value. At each iteration epsilon is reduced by 0.01, and if
                  it is not less than 0.0, another iteration is performed.
                  Thus, a value of 0.5 leads to 50 iterations. Specify a
                  smaller value if less iterations are desired.
        verboseFlag: flag to indicate whether to print some intermediate values
            or not. Specify 'False' if you want no message printed.
62 63

    Outputs:
Amir Mohammadi's avatar
Amir Mohammadi committed
64 65 66 67 68 69
        sfi: numpy array of shape (3, maxY, maxX), containing speckle-free
            image
        diffuse: numpy array of shape (3, maxY, maxX), containing the diffuse
            component
        speckleResidue: numpy array of shape (3, maxY, maxX), containing
            specular component.
70
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
71 72 73 74 75
    # initialize resulting diffuse-image
    diffuse = np.copy(srcImage)  # this copy will be updated in zIteration()

    srcPixelStatus, sfi, srcMaxChroma, srcMax, srcTotal = specular_free_image(
        diffuse)
76 77
    # sfi is the specular-free image.

Amir Mohammadi's avatar
Amir Mohammadi committed
78 79 80 81 82 83 84 85
    epsilon = startEps
    step = 0.01

    while(epsilon >= 0.0):
        # if verboseFlag:
        #     print('*')
        diffuse, srcPixelStatus = iteration(
            diffuse, srcPixelStatus, sfi, epsilon, verboseFlag)
86
        epsilon -= step
Amir Mohammadi's avatar
Amir Mohammadi committed
87 88
        # if verboseFlag:
        #     print('ep: %f' % epsilon)
89 90 91

    speckleResidue = srcImage - diffuse

Amir Mohammadi's avatar
Amir Mohammadi committed
92 93 94 95 96
    # speckleResidue is 3D but for every pixel all channels have the same
    # value.
    return sfi, diffuse, speckleResidue


97 98
def specular_free_image(src, srcPixelStatus=None):
    """Generates specular-free version of input RGB color image src.
Amir Mohammadi's avatar
Amir Mohammadi committed
99 100

    Args:
101
        src: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
102 103 104 105 106
        srcPixelStatus: numpy array of shape (maxX, maxY) containing a marker
            per pixel, indicating the type of pixel. This array is updated
            inside this function, so if the input param. is None, it is first
            allocated and initialized to 0-values.

107
    Outputs:
Amir Mohammadi's avatar
Amir Mohammadi committed
108 109 110 111 112 113 114 115 116 117
        srcPixelStatus: numpy array of shape (maxX, maxY) containing a marker
            per pixel, indicating the type of pixel.
        sfi: numpy array of shape (3, maxY, maxX) containing the specular-free
            image corresponding to src.
        srcMaxChroma: numpy array of shape (maxX, maxY) containing the max.
            chroma for each pixel.
        srcMax: numpy array of shape (maxX, maxY) containing the max. among the
            3 color-intensities, for each pixel in src.
        srcTotal: numpy array of shape (maxX, maxY) containing pixel-wise sum
            of color-intensities in src.
118
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
119

120 121 122
    # allocate output image
    cLambda = 0.6       # arbitrary constant. See paper referenced above.
    camDark = 10.0      # threshold to determine if a pixel is too dark
Amir Mohammadi's avatar
Amir Mohammadi committed
123 124
    lambdaConst = 3.0 * cLambda - 1.0

125
    if srcPixelStatus is None:
Amir Mohammadi's avatar
Amir Mohammadi committed
126
        srcPixelStatus = np.zeros((src.shape[1], src.shape[2]))
127
    else:
Amir Mohammadi's avatar
Amir Mohammadi committed
128 129 130 131 132
        assert(
            src.shape[1:-1] == srcPixelStatus.shape), "specular_free_image()" \
            ":: srcPixelStatus should be None, or should match the " \
            "pixel-layout of src."

133
    srcPixelStatus[np.logical_and.reduce(src < camDark, 0)] = G_CAMERA_DARK
134 135

    srcMaxChroma, srcMax, srcTotal = max_chroma(src)
Amir Mohammadi's avatar
Amir Mohammadi committed
136 137 138 139

    numer = srcMax * (3.0 * srcMaxChroma - 1.0)
    denom = srcMaxChroma * lambdaConst
    old_settings = np.seterr(divide='ignore', invalid='ignore')
140
    dI = np.divide(numer, denom)
Amir Mohammadi's avatar
Amir Mohammadi committed
141 142 143 144 145 146 147 148 149 150 151
    np.seterr(**old_settings)
    sI = (srcTotal - dI) / 3.0

    # src: 3d array (rgb image); sI: 2D array matching pixel-layout of src
    drgb = src.astype(float) - sI

    drgb[np.where(np.isinf(drgb))] = 0
    drgb[np.where(np.isnan(drgb))] = 0
    drgb[np.where(drgb < 0)] = 0
    drgb[np.where(drgb > 255)] = 255
    sfi = drgb
152 153 154 155 156 157

    return srcPixelStatus, sfi, srcMaxChroma, srcMax, srcTotal


def iteration(src, srcPixelStatus, sfi, epsilon, verboseFlag=True):
    """Iteratively converts each specular-pixel to diffuse.
Amir Mohammadi's avatar
Amir Mohammadi committed
158 159

    Args:
160
        src: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
161 162 163 164 165
        srcPixelStatus: numpy array of shape (maxX, maxY) containing a marker
            per pixel, indicating the type of pixel.
        sfi: numpy array of shape (3, maxY, maxX) containing the speckle-free
            image corresponding to src
        epsilon: floating-point (small) value.
166
        verboseFlag: indicator to print something in every loop.
Amir Mohammadi's avatar
Amir Mohammadi committed
167

168
    Outputs:
Amir Mohammadi's avatar
Amir Mohammadi committed
169 170 171 172
        src: numpy array of shape (3, maxY, maxX) containing updated input
            image
        srcPixelStatus: numpy array of shape (maxX, maxY) containing updated
            pixel-markers
173 174
    """

Amir Mohammadi's avatar
Amir Mohammadi committed
175 176 177
    thR = 0.1
    thG = 0.1
    pcount = 0
178 179 180
    count, srcPixelStatus = init_labels(src, sfi, epsilon, srcPixelStatus)

    while(1):
181
        maxY, maxX = src.shape[-2:]
182 183
        srcMaxChroma, srcMax, srcTotal = max_chroma(src)
        srcChroma, srcTotal = rgb_chroma(src, srcTotal)
Amir Mohammadi's avatar
Amir Mohammadi committed
184 185 186 187 188 189 190

        red_chroma_diff_x = np.diff(srcChroma[0, :, :], axis=1)
        red_chroma_diff_y = np.diff(srcChroma[0, :, :], axis=0)

        grn_chroma_diff_x = np.diff(srcChroma[1, :, :], axis=1)
        grn_chroma_diff_y = np.diff(srcChroma[1, :, :], axis=0)

191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320
        # a non-complete attempt to remove the for loop below
        # # find non G_CAMERA_DARK pixels
        # non_g_camera_dark = (srcPixelStatus != G_CAMERA_DARK)
        # non_g_camera_dark[-1, :] = False
        # non_g_camera_dark[:, -1] = False
        # g_specularx = (non_g_camera_dark) & \
        #     (srcPixelStatus == G_SPECULARX)
        # g_speculary = (non_g_camera_dark) & \
        #     (srcPixelStatus == G_SPECULARY)

        # '''for specular y'''

        # # if it is a boundary in x
        # boundary_x = np.zeros_like(g_specularx)
        # boundary_x[:-1, :-1] = g_specularx[:-1, :-1] & \
        #     (np.fabs(red_chroma_diff_x[:-1, :]) > thR) & \
        #     (np.fabs(grn_chroma_diff_x[:-1, :]) > thG)
        # srcPixelStatus[boundary_x] = G_BOUNDARY

        # # if it is noise in x
        # srcMaxChromaDiffX = np.diff(srcMaxChroma, axis=1)
        # srcMaxChromaAbsDiffX = np.fabs(srcMaxChromaDiffX)
        # noise_mask_x = g_specularx & (~boundary_x)
        # noise_mask_x[:-1, :-1] = noise_mask_x[:-1, :-1] & \
        #     (srcMaxChromaAbsDiffX[:-1, :] < 0.01)
        # srcPixelStatus[noise_mask_x] = G_NOISE

        # # if it is a boundary in x
        # boundary_y = np.zeros_like(g_speculary)
        # boundary_y[:-1, :-1] = g_speculary[:-1, :-1] & \
        #     (np.fabs(red_chroma_diff_y[:, :-1]) > thR) & \
        #     (np.fabs(grn_chroma_diff_y[:, :-1]) > thG)
        # srcPixelStatus[boundary_y] = G_BOUNDARY

        # # if it is noise in x
        # srcMaxChromaDiffY = np.diff(srcMaxChroma, axis=0)
        # srcMaxChromaAbsDiffY = np.fabs(srcMaxChromaDiffY)
        # noise_mask_y = g_speculary & (~boundary_y)
        # noise_mask_y[:-1, :-1] = noise_mask_y[:-1, :-1] & \
        #     (srcMaxChromaAbsDiffY[:, :-1] < 0.01)
        # srcPixelStatus[noise_mask_y] = G_NOISE

        # # if it is not noise or boundary
        # boundary_x = np.zeros_like(g_specularx)
        # boundary_x[:-1, :-1] = g_specularx[:-1, :-1] & \
        #     (np.fabs(red_chroma_diff_x[:-1, :]) > thR) & \
        #     (np.fabs(grn_chroma_diff_x[:-1, :]) > thG)
        # srcMaxChromaDiffX = np.diff(srcMaxChroma, axis=1)
        # srcMaxChromaAbsDiffX = np.fabs(srcMaxChromaDiffX)
        # noise_mask_x = g_specularx & (~boundary_x)
        # noise_mask_x[:-1, :-1] = noise_mask_x[:-1, :-1] & \
        #     (srcMaxChromaAbsDiffX[:-1, :] < 0.01)
        # non_noise_mask_x = g_specularx & (~boundary_x) & (~noise_mask_x)
        # srcPixelStatus[non_noise_mask_x] = G_DIFFUSE
        # non_noise_mask_next_x = np.roll(non_noise_mask_x, 1, axis=1)
        # srcPixelStatus[non_noise_mask_next_x] = G_DIFFUSE

        # boundary_y = np.zeros_like(g_speculary)
        # boundary_y[:-1, :-1] = g_speculary[:-1, :-1] & \
        #     (np.fabs(red_chroma_diff_y[:, :-1]) > thR) & \
        #     (np.fabs(grn_chroma_diff_y[:, :-1]) > thG)
        # srcMaxChromaDiffY = np.diff(srcMaxChroma, axis=0)
        # srcMaxChromaAbsDiffY = np.fabs(srcMaxChromaDiffY)
        # noise_mask_y = g_speculary & (~boundary_y)
        # noise_mask_y[:-1, :-1] = noise_mask_y[:-1, :-1] & \
        #     (srcMaxChromaAbsDiffY[:, :-1] < 0.01)
        # non_noise_mask_y = g_speculary & (~boundary_y) & (~noise_mask_y)
        # srcPixelStatus[non_noise_mask_y] = G_DIFFUSE
        # non_noise_mask_next_y = np.roll(non_noise_mask_y, 1, axis=0)
        # srcPixelStatus[non_noise_mask_next_y] = G_DIFFUSE

        # # if it is not boundary or noise
        # pos_chroma_maskx = (srcMaxChromaDiffX[:-1, :] > 0)
        # reduce_mask_pos_x = non_noise_mask_x.copy()
        # reduce_mask_pos_x[:-1, :-1] = reduce_mask_pos_x[:-1, :-1] & \
        #     pos_chroma_maskx
        # reduce_mask_pos_next_x = np.roll(reduce_mask_pos_x, 1, axis=1)
        # iro = np.moveaxis(np.moveaxis(src, 0, -1)[reduce_mask_pos_x], -1, 0)
        # refMaxChroma = srcMaxChroma[reduce_mask_pos_next_x]
        # iroTotal = srcTotal[reduce_mask_pos_x]
        # iroMax = srcMax[reduce_mask_pos_x]
        # iroMaxChroma = srcMaxChroma[reduce_mask_pos_x]
        # pStat, iroRes = specular_to_diffuse_vectorized(
        #     iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)
        # iro[:, pStat != G_NOISE] = iroRes[:, pStat != G_NOISE]

        # pos_chroma_maskx = ~pos_chroma_maskx
        # reduce_mask_pos_x = non_noise_mask_x.copy()
        # reduce_mask_pos_x[:-1, :-1] = reduce_mask_pos_x[:-1, :-1] & \
        #     pos_chroma_maskx
        # reduce_mask_pos_next_x = np.roll(reduce_mask_pos_x, 1, axis=1)
        # iro = np.moveaxis(np.moveaxis(src, 0, -1)[reduce_mask_pos_next_x],
        #                   -1, 0)
        # refMaxChroma = srcMaxChroma[reduce_mask_pos_x]
        # iroTotal = srcTotal[reduce_mask_pos_next_x]
        # iroMax = srcMax[reduce_mask_pos_next_x]
        # iroMaxChroma = srcMaxChroma[reduce_mask_pos_next_x]
        # pStat, iroRes = specular_to_diffuse_vectorized(
        #     iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)
        # iro[:, pStat != G_NOISE] = iroRes[:, pStat != G_NOISE]

        # # if it is not boundary or noise
        # pos_chroma_masky = (srcMaxChromaDiffY[:, :-1] > 0)
        # reduce_mask_pos_y = non_noise_mask_y.copy()
        # reduce_mask_pos_y[:-1, :-1] = reduce_mask_pos_y[:-1, :-1] & \
        #     pos_chroma_masky
        # reduce_mask_pos_next_y = np.roll(reduce_mask_pos_y, 1, axis=1)
        # iro = np.moveaxis(np.moveaxis(src, 0, -1)[reduce_mask_pos_y], -1, 0)
        # refMaxChroma = srcMaxChroma[reduce_mask_pos_next_y]
        # iroTotal = srcTotal[reduce_mask_pos_y]
        # iroMax = srcMax[reduce_mask_pos_y]
        # iroMaxChroma = srcMaxChroma[reduce_mask_pos_y]
        # pStat, iroRes = specular_to_diffuse_vectorized(
        #     iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)
        # iro[:, pStat != G_NOISE] = iroRes[:, pStat != G_NOISE]

        # pos_chroma_masky = ~pos_chroma_masky
        # reduce_mask_pos_y = non_noise_mask_y.copy()
        # reduce_mask_pos_y[:-1, :-1] = reduce_mask_pos_y[:-1, :-1] & \
        #     pos_chroma_masky
        # reduce_mask_pos_next_y = np.roll(reduce_mask_pos_y, 1, axis=1)
        # iro = np.moveaxis(np.moveaxis(src, 0, -1)[reduce_mask_pos_next_y],
        #                   -1, 0)
        # refMaxChroma = srcMaxChroma[reduce_mask_pos_y]
        # iroTotal = srcTotal[reduce_mask_pos_next_y]
        # iroMax = srcMax[reduce_mask_pos_next_y]
        # iroMaxChroma = srcMaxChroma[reduce_mask_pos_next_y]
        # pStat, iroRes = specular_to_diffuse_vectorized(
        #     iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)
        # iro[:, pStat != G_NOISE] = iroRes[:, pStat != G_NOISE]
Amir Mohammadi's avatar
Amir Mohammadi committed
321 322 323
        for y in range(maxY - 1):
            for x in range(maxX - 1):
                current_pixel = srcPixelStatus[y, x]
324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432
                if(current_pixel == G_CAMERA_DARK):
                    continue

                drx = red_chroma_diff_x[y, x]
                dgx = grn_chroma_diff_x[y, x]
                dry = red_chroma_diff_y[y, x]
                dgy = grn_chroma_diff_y[y, x]

                if(current_pixel == G_SPECULARX):
                    # if it is  a boundary in the x direction
                    if (abs(drx) > thR) and (abs(dgx) > thG):
                        # pixel right
                        srcPixelStatus[y, x] = G_BOUNDARY
                        continue
                    # if it is a noise
                    if abs(srcMaxChroma[y, x] - srcMaxChroma[y, x + 1]) < 0.01:
                        srcPixelStatus[y, x] = G_NOISE
                        continue
                    # reduce the specularity at x direction
                    if srcMaxChroma[y, x] < srcMaxChroma[y, x + 1]:
                        iro = src[:, y, x]
                        # pStat = current_pixel
                        refMaxChroma = srcMaxChroma[y, x + 1]
                        iroTotal = srcTotal[y, x]
                        iroMax = srcMax[y, x]
                        iroMaxChroma = srcMaxChroma[y, x]
                        pStat, iroRes = specular_to_diffuse(
                            iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)

                        # update pixelStatus only if
                        # zSpecular2Diffuse() returned pStat as
                        # G_NOISE.
                        if pStat != G_NOISE:
                            src[:, y, x] = iroRes
                        # else:
                        #     srcPixelStatus[y, x] = pStat

                    else:
                        iro = src[:, y, x + 1]
                        # pStat = srcPixelStatus[y, x + 1]
                        refMaxChroma = srcMaxChroma[y, x]
                        iroTotal = srcTotal[y, x + 1]
                        iroMax = srcMax[y, x + 1]
                        iroMaxChroma = srcMaxChroma[y, x + 1]
                        pStat, iroRes = specular_to_diffuse(
                            iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)

                        # update pixelStatus only if
                        # zSpecular2Diffuse() returned pStat as
                        # G_NOISE.
                        if pStat != G_NOISE:
                            src[:, y, x + 1] = iroRes
                        # else:
                        #     srcPixelStatus[y, x + 1] = pStat

                    srcPixelStatus[y, x] = G_DIFFUSE
                    srcPixelStatus[y, x + 1] = G_DIFFUSE

                if(current_pixel == G_SPECULARY):
                    # if it is a boundary in the y direction
                    if (abs(dry) > thR) and (abs(dgy) > thG):
                        # pixel right
                        srcPixelStatus[y, x] = G_BOUNDARY
                        continue
                    # if it is a noise
                    if abs(srcMaxChroma[y, x] - srcMaxChroma[y + 1, x]) < 0.01:
                        srcPixelStatus[y, x] = G_NOISE
                        continue
                    # reduce the specularity in y direction
                    if srcMaxChroma[y, x] < srcMaxChroma[y + 1, x]:
                        iro = src[:, y, x]
                        pStat = current_pixel
                        refMaxChroma = srcMaxChroma[y + 1, x]
                        iroTotal = srcTotal[y, x]
                        iroMax = srcMax[y, x]
                        iroMaxChroma = srcMaxChroma[y, x]
                        pStat, iroRes = specular_to_diffuse(
                            iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma)
                        # C call:
                        # zSpecular2Diffuse(src(y,x),src(y+1,x).zMaxChroma())

                        # update pixelStatus only if
                        # zSpecular2Diffuse() returned pStat as
                        # G_NOISE.
                        if pStat == G_NOISE:
                            srcPixelStatus[y, x] = pStat
                        else:
                            src[:, y, x] = iroRes

                    else:
                        iro = src[:, y + 1, x]
                        pStat = srcPixelStatus[y + 1, x]
                        pMaxChroma = srcMaxChroma[y + 1, x]
                        iroTotal = srcTotal[y + 1, x]
                        iroMax = srcMax[y + 1, x]
                        iroMaxChroma = srcMaxChroma[y + 1, x]
                        pStat, iroRes = specular_to_diffuse(
                            iro, iroMax, iroMaxChroma, iroTotal, pMaxChroma)

                        # update pixelStatus only if
                        # zSpecular2Diffuse() returned pStat as
                        # G_NOISE.
                        if pStat == G_NOISE:
                            srcPixelStatus[y + 1, x] = pStat
                        else:
                            src[:, y + 1, x] = iroRes

                    srcPixelStatus[y, x] = G_DIFFUSE
                    srcPixelStatus[y + 1, x] = G_DIFFUSE
Amir Mohammadi's avatar
Amir Mohammadi committed
433 434

        pcount = count
435 436
        count, srcPixelStatus = init_labels(src, sfi, epsilon, srcPixelStatus)

Amir Mohammadi's avatar
Amir Mohammadi committed
437 438 439
        # Robby Tan's original C++ code checks if count<0, but that doesn't
        # make sense as count cannot be negative.
        if(count == 0 or pcount <= count):
440
            break       # break out of the while-loop
Amir Mohammadi's avatar
Amir Mohammadi committed
441

442
    srcPixelStatus = reset_labels(srcPixelStatus)
Amir Mohammadi's avatar
Amir Mohammadi committed
443

444 445 446 447 448
    return src, srcPixelStatus


def init_labels(src, sfi, epsilon, srcPixelStatus):
    """Generates initial labels for all pixels in src (input RGB image).
Amir Mohammadi's avatar
Amir Mohammadi committed
449 450

    Args:
451
        src: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
452 453
        sfi: numpy array of shape (3, maxY, maxX) containing a speckle-free
            image corresponding to src.
454
        epsilon: positive floating point (small) value
Amir Mohammadi's avatar
Amir Mohammadi committed
455 456 457
        srcPixelStatus: numpy array of shape (maxX, maxY) containing pixel-
            markers corresponding to src.

458
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
459 460 461
        count: number of pixels marked as specular in src.
        srcPixelStatus: numpy array of shape (maxX, maxY) containing updated
            pixel-markers corresponding to src.
462
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
463 464 465

    old_settings = np.seterr(divide='ignore', invalid='ignore')
    # to have initial labels
466 467 468 469 470 471 472 473
    zTotalSrc = np.sum(src, axis=0)
    diff_x_src = np.diff(zTotalSrc, axis=1)
    diff_y_src = np.diff(zTotalSrc, axis=0)

    zTotalSfi = np.sum(sfi, axis=0)
    diff_x_sfi = np.diff(zTotalSfi, axis=1)
    diff_y_sfi = np.diff(zTotalSfi, axis=0)

474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496
    dlog_src_x = np.log(np.fabs(diff_x_src))
    dlog_src_y = np.log(np.fabs(diff_y_src))

    dlog_sfi_x = np.log(np.fabs(diff_x_sfi))
    dlog_sfi_y = np.log(np.fabs(diff_y_sfi))

    dlogx = np.fabs(dlog_src_x - dlog_sfi_x)
    dlogy = np.fabs(dlog_src_y - dlog_sfi_y)

    # maxY = src.shape[1]
    # maxX = src.shape[2]
    valid_values = (G_BOUNDARY, G_NOISE, G_CAMERA_DARK, G_DIFFUSE)
    notvalid_mask = np.in1d(
        srcPixelStatus[1:-1, 1:-1], valid_values,
        invert=True).reshape(srcPixelStatus.shape[0] - 2, -1)
    dlogx_mask = (notvalid_mask) & (dlogx[1:-1, 1:] > epsilon)
    dlogy_mask = (notvalid_mask) & (dlogy[1:, 1:-1] > epsilon) & \
        (~(dlogx_mask))
    diffuse_mask = (notvalid_mask) & (~dlogx_mask) & (~dlogy_mask)
    count = notvalid_mask.sum()
    srcPixelStatus[1:-1, 1:-1][dlogx_mask] = G_SPECULARX
    srcPixelStatus[1:-1, 1:-1][dlogy_mask] = G_SPECULARY
    srcPixelStatus[1:-1, 1:-1][diffuse_mask] = G_DIFFUSE
Amir Mohammadi's avatar
Amir Mohammadi committed
497 498

    np.seterr(**old_settings)
499 500 501 502
    return count, srcPixelStatus    # count is the number of specular pixels


def specular_to_diffuse(iro, iroMax, iroMaxChroma, iroTotal, refMaxChroma):
Amir Mohammadi's avatar
Amir Mohammadi committed
503 504 505 506 507 508
    """Converts a color-pixel from speckled to diffuse, by subtracting a
    certain amount from its intensity value in each color channel.

    Args:
        iro: 3-element column vector containing the rgb-color values of the
            pixel to be processed.
509 510 511 512
        iroMax: max value among the elements of iro
        iroMaxChroma: max-chroma for this pixel
        iroTotal: sum of the 3 elements of iro.
        refMaxChroma: chroma-value of a neighboring pixel for comparison
Amir Mohammadi's avatar
Amir Mohammadi committed
513

514
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
515 516
        pixelStatus: a value (marker) indicating whether the pixel is
            considered as noise or diffuse, after processing.
517 518
        iro: updated pixel-color-values.
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
519 520 521 522 523
    c = iroMaxChroma
    pixelStatus = 0  # arbitrary initialization
    numr = (iroMax * (3.0 * c - 1.0))
    denm = (c * (3.0 * refMaxChroma - 1.0))
    if abs(denm) > 0.000000001:  # to avoid div. by zero.
524
        dI = numr / denm
Amir Mohammadi's avatar
Amir Mohammadi committed
525
        sI = (iroTotal - dI) / 3.0
526
        nrgb = iro - sI
Amir Mohammadi's avatar
Amir Mohammadi committed
527 528

        if np.any(nrgb < 0):
529 530 531 532 533 534 535 536
            pixelStatus = G_NOISE
        else:
            iro = nrgb
    else:
        pixelStatus = G_NOISE

    return pixelStatus, iro

Amir Mohammadi's avatar
Amir Mohammadi committed
537

Amir Mohammadi's avatar
Amir Mohammadi committed
538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572
# def specular_to_diffuse_vectorized(iro, iroMax, iroMaxChroma, iroTotal,
#                                    refMaxChroma):
#     """Converts a color-pixel from speckled to diffuse, by subtracting a
#     certain amount from its intensity value in each color channel.
#
#     Args:
#         iro: 3-element column vector containing the rgb-color values of the
#             pixel to be processed.
#         iroMax: max value among the elements of iro
#         iroMaxChroma: max-chroma for this pixel
#         iroTotal: sum of the 3 elements of iro.
#         refMaxChroma: chroma-value of a neighboring pixel for comparison
#
#     Returns:
#         pixelStatus: a value (marker) indicating whether the pixel is
#             considered as noise or diffuse, after processing.
#         iro: updated pixel-color-values.
#     """
#     c = iroMaxChroma
#     # arbitrary initialization
#     pixelStatus = np.empty((*iro.shape[1:]), dtype=iro.dtype)
#     numr = (iroMax * (3.0 * c - 1.0))
#     denm = (c * (3.0 * refMaxChroma - 1.0))
#
#     non_zero = np.fabs(denm) > 0.000000001
#     dI = numr[non_zero] / denm[non_zero]
#     sI = (iroTotal[non_zero] - dI) / 3.0
#     nrgb = iro[:, non_zero] - sI
#     negnrgb = np.logical_or.reduce(nrgb < 0, axis=0)
#
#     pixelStatus[non_zero][negnrgb] = G_NOISE
#     iro[:, non_zero][:, ~negnrgb] = nrgb[:, ~negnrgb]
#     pixelStatus[~non_zero] = G_NOISE
#
#     return pixelStatus, iro
Amir Mohammadi's avatar
Amir Mohammadi committed
573 574


575 576
def reset_labels(pixelLabels):
    """Resets all pixel-markers (except those indicating "dark" pixels, to 0.
577
    Called in zIteration().
Amir Mohammadi's avatar
Amir Mohammadi committed
578 579 580 581 582

    Args:
        pixelLabels: numpy array of shape (maxX, maxY) containing pixel-
            markers.

583
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
584 585
        pixelLabels: numpy array of shape (maxX, maxY) containing updated
            pixel-markers..
586 587
    """
    # to reset the label of the pixels
588
    pixelLabels[np.where(pixelLabels != G_CAMERA_DARK)] = 0
589 590 591 592 593

    return pixelLabels


#####
Amir Mohammadi's avatar
Amir Mohammadi committed
594
# functions from globals
595 596
#####

Amir Mohammadi's avatar
Amir Mohammadi committed
597

598
def max_color(rgbImg):
Amir Mohammadi's avatar
Amir Mohammadi committed
599 600 601 602
    """For input RGB image, this function computes max intensity among all
    color-channels, for each pixel position.

    Args:
603
        rgbImg: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
604

605
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
606 607
        rgbMax: (2D numpy array of shape (maxY, maxY)) contains pixel-wise max.
            among the three color values.
608
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
609

610 611 612 613
    return np.amax(rgbImg, axis=0)


def total_color(rgbImg):
Amir Mohammadi's avatar
Amir Mohammadi committed
614 615 616 617
    """For input RGB image, this function computes sum of intensities in each
    color-channel for each pixel position.

    Args:
618
        rgbImg: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
619

620
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
621 622
        rgbTotal: (2D numpy array of shape (maxY, maxY)) contains pixel-wise
            sum of the 3 color-values.
623
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
624

625 626
    return np.sum(rgbImg.astype(float), axis=0)

Amir Mohammadi's avatar
Amir Mohammadi committed
627

628
def max_chroma(rgbImg, rgbMax=None, rgbTotal=None):
Amir Mohammadi's avatar
Amir Mohammadi committed
629 630 631 632
    """Given an input RGB image (rgbImg, with shape (3, maxY, maxX)), this
    function computes the maximum chroma-value for each pixel position.

    Args:
633
        rgbImg: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
634 635 636 637 638
        rgbMax: numpy array of shape (maxY, maxX) containing pixel-wise max.
            color intensity. If None, it is computed.
        rgbTotal: numpy array of shape (maxY, maxX) containing pixel-wise sum
            of color-intensities. If None, it is computed.

639
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
640 641 642 643 644 645
        rgbChroma: (3D numpy array of same shape as rgbImg) contains the
            chroma-values of the individual channels in rgbChroma, and
        rgbMax: (2D numpy array of shape (maxY, maxY)) contains pixel-wise max.
            among the three color values.
        rgbTotal: (2D numpy array of shape (maxY, maxY)) contains pixel-wise
            sum of the 3 color-values.
646
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
647

648 649 650 651
    if rgbMax is None:
        rgbMax = max_color(rgbImg)
    if rgbTotal is None:
        rgbTotal = total_color(rgbImg)
Amir Mohammadi's avatar
Amir Mohammadi committed
652 653

    old_settings = np.seterr(divide='ignore', invalid='ignore')
654
    maxChroma = np.divide(rgbMax.astype(float), rgbTotal.astype(float))
Amir Mohammadi's avatar
Amir Mohammadi committed
655 656 657 658 659 660
    np.seterr(**old_settings)
    maxChroma[np.where(rgbTotal == 0)] = 0

    # max_chroma(rgbImg, rgbMax=None, rgbTotal=None)
    return maxChroma, rgbMax, rgbTotal

661 662

def rgb_chroma(rgbImg, rgbTotal=None):
Amir Mohammadi's avatar
Amir Mohammadi committed
663 664 665 666
    """Given an input RGB color image, compute the chroma-values in the 3
    channels.

    Args:
667
        rgbImg: numpy array of shape (3, maxY, maxX) containing a RGB image.
Amir Mohammadi's avatar
Amir Mohammadi committed
668 669 670
        rgbTotal: numpy array of shape (maxY, maxX) containing pixel-wise sum
            of color-intensities. If None, it is computed.

671
    Returns:
Amir Mohammadi's avatar
Amir Mohammadi committed
672 673 674 675
        rgbChroma: (3D numpy array of same shape as rgbImg) contains the
            chroma-values of the individual channels in rgbChroma, and
        rgbTotal: (2D numpy array of shape (maxY, maxY)) contains pixel-wise
            sum of the 3 color-values.
676
    """
Amir Mohammadi's avatar
Amir Mohammadi committed
677

678 679
    if rgbTotal is None:
        rgbTotal = total_color(rgbImg)
Amir Mohammadi's avatar
Amir Mohammadi committed
680 681

    old_settings = np.seterr(divide='ignore', invalid='ignore')
682
    rgbChroma = np.divide(rgbImg.astype(float), rgbTotal.astype(float))
Amir Mohammadi's avatar
Amir Mohammadi committed
683
    np.seterr(**old_settings)
684
    rgbChroma[:, rgbTotal == 0] = 0
685 686 687 688 689 690 691

    return rgbChroma, rgbTotal


'''
utility function to display image on screen (for debugging).
'''
Amir Mohammadi's avatar
Amir Mohammadi committed
692 693


694
def imshow(image):
Amir Mohammadi's avatar
Amir Mohammadi committed
695
    '''please use bob.io.image.imshow instead'''
696 697
    import matplotlib
    from matplotlib import pyplot as plt
Amir Mohammadi's avatar
Amir Mohammadi committed
698 699 700
    if len(image.shape) == 3:
        # imshow() expects color image in a slightly different format, so first
        # rearrange the 3d data for imshow...
701 702 703
        outImg = image.tolist()
        result = np.dstack((outImg[0], outImg[1]))
        outImg = np.dstack((result, outImg[2]))
Amir Mohammadi's avatar
Amir Mohammadi committed
704 705 706
        # [:,:,1], cmap=mpl.cm.gray)
        plt.imshow((outImg * 255.0).astype(np.uint8))

707
    else:
Amir Mohammadi's avatar
Amir Mohammadi committed
708 709
        if(len(image.shape) == 2):
            # display gray image.
710 711
            plt.imshow(image.astype(np.uint8), cmap=matplotlib.cm.gray)
        else:
712
            print("inshow():: image should be either 2d or 3d.")
Amir Mohammadi's avatar
Amir Mohammadi committed
713

714 715
    plt.show()

Amir Mohammadi's avatar
Amir Mohammadi committed
716

717
def computeIQSpecularityFeatures(rgbImage, startEps=0.05):
Amir Mohammadi's avatar
Amir Mohammadi committed
718 719 720 721 722

    speckleFreeImg, diffuseImg, speckleImg = remove_highlights(
        rgbImage.astype(float), startEps=0.05)

    if len(speckleImg.shape) == 3:
723
        speckleImg = speckleImg[0]
Amir Mohammadi's avatar
Amir Mohammadi committed
724

725 726 727
    speckleImg = speckleImg.clip(min=0)

    speckleMean = np.mean(speckleImg)
Amir Mohammadi's avatar
Amir Mohammadi committed
728 729 730 731 732 733 734 735 736 737
    lowSpeckleThresh = speckleMean * 1.5
    hiSpeckleThresh = speckleMean * 4.0
    specklePixels = speckleImg[
        np.where(
            np.logical_and(
                speckleImg > lowSpeckleThresh,
                speckleImg < hiSpeckleThresh))]

    r = float(specklePixels.shape[0]) / \
        (speckleImg.shape[0] * speckleImg.shape[1])
738
    m = np.mean(specklePixels)
Amir Mohammadi's avatar
Amir Mohammadi committed
739 740 741 742
    s = np.std(specklePixels)

    return (r, m / 150.0, s / 150.0)

743 744

'''
Amir Mohammadi's avatar
Amir Mohammadi committed
745 746
utility function to test the implementation. Relies on bob to load and save
images.
747
'''
Amir Mohammadi's avatar
Amir Mohammadi committed
748 749


750 751 752
def test_highlight_detection():
    inputImageFile = '/idiap/home/sbhatta/work/Antispoofing/ImageQualityMeasures/specular_highlights/highlight_C_Code/images/head.ppm'
    outRoot = '/idiap/home/sbhatta/work/Antispoofing/ImageQualityMeasures/specular_highlights/'
Amir Mohammadi's avatar
Amir Mohammadi committed
753 754 755 756 757
    # inputImageFile =
    # 'C:/IDIAP/AntiSpoofing/ImageQualityMeasures/highlight/images/head.ppm'
    # outRoot = 'C:/IDIAP/AntiSpoofing/ImageQualityMeasures/'

    # load color image
758
    inpImage = bob.io.base.load(inputImageFile)
Amir Mohammadi's avatar
Amir Mohammadi committed
759 760 761 762 763
    # inpImage = misc.imread(inputImageFile)
    # inpImage =  np.rollaxis(inpImage,2)

    speckleFreeImg, diffuseImg, speckleImg = remove_highlights(
        inpImage.astype(float), startEps=0.05)
764

765
    print('saving output images')
Amir Mohammadi's avatar
Amir Mohammadi committed
766 767 768 769 770 771 772 773 774 775
    bob.io.base.save(
        speckleFreeImg.astype('uint8'),
        outRoot + 'speckleFreeHeadImage.ppm')
    # speckleFreeImg = np.rollaxis(np.rollaxis(speckleFreeImg, 0,-1),2,1)
    # misc.imsave(outRoot+'speckleFreeHeadImage.ppm', speckleFreeImg.astype('uint8'))

    bob.io.base.save(
        diffuseImg.astype('uint8'),
        outRoot +
        'diffuseImgHeadImage.ppm')
776 777
#     diffuseImg = np.rollaxis(np.rollaxis(diffuseImg, 0,-1),2,1)
#     misc.imsave(outRoot+'diffuseImgHeadImage.ppm', diffuseImg.astype('uint8'))
Amir Mohammadi's avatar
Amir Mohammadi committed
778 779 780 781 782

    bob.io.base.save(
        speckleImg.astype('uint8'),
        outRoot +
        'speckleImgHeadImage.ppm')
783 784
#     speckleImg = np.rollaxis(np.rollaxis(speckleImg, 0,-1),2,1)
#     misc.imsave(outRoot+'speckleImgHeadImage.ppm', speckleImg.astype('uint8'))
Amir Mohammadi's avatar
Amir Mohammadi committed
785 786
#

787
#     r,m,s=computeIQSpecularityFeatures(inpImage, startEps=0.05)
788
#     print(r, m, s)
789

Amir Mohammadi's avatar
Amir Mohammadi committed
790
#
791 792 793 794 795
#    imshow(inpImage)
#    imshow(speckleFreeImg)
#    imshow(diffuseImg)
    imshow(speckleImg)

Amir Mohammadi's avatar
Amir Mohammadi committed
796

797
if __name__ == '__main__':
798
    test_highlight_detection()