Adding pixel-level classification to ee.ImageCollection

Question

I am working with ee.ImageCollection('LANDSAT/LE09/C02/T1_L2') data pulled from earthengine-api in Python.

I have previously defined a function to mask cloud, shadow, water, and snow. based on the bit properties of the QA_PIXEL band, as so:

def mask_L9(image):
    water = 1 << 7
    shadow = 1 << 4
    snow = 1 << 5
    cloud = 1 << 3

    # Get the QA band.
    qa = image.select('QA_PIXEL')

    # Flags should be set to zero.
    mask = qa.bitwiseAnd(water).eq(0).And(
           qa.bitwiseAnd(shadow).eq(0).And(
           qa.bitwiseAnd(snow).eq(0).And(
           qa.bitwiseAnd(cloud).eq(0))))

    return image.select(['SR_B.*']).multiply(0.0000275).add(-0.2) \
        .updateMask(mask) \
        .set('system:time_start', image.get('system:time_start'))

l9 = (ee.ImageCollection('LANDSAT/LC09/C02/T1_L2').filterDate('2021-10-31', '2023-09-30')
                                                  .map(mask_L9)

I would like to modify this function, so that instead of masking cloud, shadow, water, and snow, a new band is added to the images that specifies their classification. The modified function I have written is:

def addPixelClass_L9(image):
    water = 1 << 7
    shadow = 1 << 4
    snow = 1 << 5
    cloud = 1 << 3
    
    # Get the QA band.
    qa = image.select('QA_PIXEL')

    class = 0

    if qa.bitwiseAnd(water_mask).neq(0):
        class = 7
    elif qa.bitwiseAnd(cloud_shadow_mask).neq(0):
        class = 4
    elif qa.bitwiseAnd(snow_mask).neq(0):
        class = 5
    elif qa.bitwiseAnd(cloud_mask).neq(0):
        class = 3
    else:
        class = 0
    
    return image.select(['SR_B.*']).multiply(0.0000275).add(-0.2) \
           .addBands(ee.Image.constant(class).rename('CLASS')) \
           .set('system:time_start', image.get('system:time_start'))

However, this returns all pixels as class 7 (i.e., water), which I know isn't the case. Can anyone provide advice as to where I am going wrong? I don't have much experience working in Python, so am unsure if this is something simple re: syntax that I am missing. Cheers.

Answer 1

I believe I found the answer in the end by using .where instead of an if/else expression:

def addPixelClass_L9(image):
    water = 1 << 7
    shadow = 1 << 4
    snow = 1 << 5
    cloud = 1 << 3

    # Get the QA band.
    qa = image.select('QA_PIXEL')

    class_value = ee.Image.constant(1)
    class_value = class_value.where(qa.bitwiseAnd(water), 7)
    class_value = class_value.where(qa.bitwiseAnd(shadow), 4)
    class_value = class_value.where(qa.bitwiseAnd(snow), 5)
    class_value = class_value.where(qa.bitwiseAnd(cloud), 3)

    return image.select(['SR_B.*']).multiply(0.0000275).add(-0.2) \
           .addBands(class_value.rename('qa_classification')) \
           .set('system:time_start', image.get('system:time_start'))