import OpenEXR
import cv2
import numpy as np
import imageio
import ast
import os
import json
import codecs
import csv

exr_filename = "path_to_your_exr.exr"

# Load the exr file into memory
exr_file = OpenEXR.File(exr_filename)

# the exr header contains usefull information about the EXR file
exr_header = exr_file.header()
print(exr_header)

# The manifest gives information about the Cryptomatte data.
# If you dont need cryptomatte skip it
# the '3ae39a5' may be different for you, but you can see it in the printed header
exr_manifests = ast.literal_eval(exr_header['cryptomatte/3ae39a5/manifest']) 
print(exr_manifests)

# The parts contains all the layers we need
# ALL LAYER NAME COMES FROM THIS
exr_parts = exr_file.parts
print(exr_parts)



# GET DEPTH
depth_part = next((part for part in exr_parts if part.name() == "depth"), None)
depth_pixels = depth_part.channels['depth.z'].pixels

depth_image_filename = "path_to_your_depth_file.tif"
imageio.imwrite(depth_image_filename, depth_pixels)



# GET NORMAL
cameranormal_part = next((part for part in exr_parts if part.name() == "N_camera"), None)
cameranormal_pixels = cameranormal_part.channels['N_camera'].pixels

cameranormal_filename = "path to your normal file.tif"
imageio.imwrite(cameranormal_filename, cameranormal_pixels)


# GET COLOR
color_part = next((part for part in exr_parts if part.name() == "C"), None)
color_file_name = "path to your color file.tif"
color_pixels =color_part.channels['RGBA'].pixels
imageio.imwrite(color_file_name, color_pixels)


# change gamma for color
postprocess_pixels = color_pixels.copy().astype(np.float32)
gamma = 2.2
postprocess_pixels = np.power(postprocess_pixels, 1.0 / gamma)

# Save postprocessed image as 8 bit RGB

postprocess_pixels_to_save = postprocess_pixels.copy()

if postprocess_pixels_to_save.ndim == 3:
    postprocess_pixels_to_save = postprocess_pixels_to_save[:,:,:3]
postprocess_pixels_to_save = np.clip(postprocess_pixels_to_save, 0, 1)
postprocess_pixels_to_save = postprocess_pixels_to_save * 255
postprocess_pixels_to_save = postprocess_pixels_to_save.astype(np.uint8)

# save postprocess
postprocess_file_name = "path to postprocessed file.png"
imageio.imwrite(postprocess_file_name, postprocess_pixels_to_save)




# Get the cryptomatte part
# I only used it for bit mask!!!!!!!!!!!!!!!

cryptomatte_part = next((part for part in exr_parts if part.name() == "CryptoObject00"), None)
cryptomatte_pixels = cryptomatte_part.channels['CryptoObject00'].pixels
#only the red channel matters
cryptomatte_pixels = cryptomatte_pixels[:, :, 0]  


# Extract segmentation from cryptomatte
# change the 'xxxxxxxxxxxxxxxxx' to a grayscale value for every cryptomatte name

cryptomatte_segmentation = np.zeros((1024,1024), dtype=np.uint8)
for key, hex_value in exr_manifests.items():
    # Convert hex string to integer
    target_int = int(hex_value, 16)

    # View the float32 array as uint32 for direct comparison
    data_uint32 = cryptomatte_pixels.view(np.uint32)

    # Create a boolean mask where the uint32 values match the target integer
    cryptomatte_mask = data_uint32 == target_int
    cryptomatte_mask = cryptomatte_mask.astype(np.uint8) * 255

    cryptomatte_mask = cryptomatte_mask != 0
    cryptomatte_segmentation[cryptomatte_mask] = xxxxxxxxxxxxxxxxx


# Save Cryptomatte segmentation
cryptomatte_file_name = "path to your semantic segmentation mask.png"
imageio.imwrite(cryptomatte_file_name, cryptomatte_segmentation)