Oiio Multilayer to Multipart with bit depth conversion

convert_to_multipart.py

import OpenImageIO as oiio
from OpenImageIO import ImageInput, ImageOutput
from OpenImageIO import ImageBuf, ImageSpec, ImageBufAlgo
from collections import OrderedDict
import fnmatch

import re
import os

path = "in.exr"

channel_rename_map = {"ViewLayer.":""}
float_layers = ["crypto", "depth", "position"]
re_frames = re.compile(r"(?:.*)([.|_][0-9]{2,10}[.|_])(?:.*)")

def construct_out_path(path):
    out_path = ".".join(path.split(".")[:-1])
    out_path = out_path + "_mp." + path.split(".")[-1]
    frames = re.search(re_frames, path)    
    if frames is not None:
        out_path = path.split(frames.groups()[0])[0]
        out_path = out_path + "_mp" + frames.groups()[0] + path.split(frames.groups()[0])[-1]
    print(out_path)
    return out_path

def is_multipart(path):
    inp = ImageInput.open(path)
    if inp:
        check = inp.seek_subimage(1, 0)
        inp.close()
        return check
    inp.close()
    return None
    
def extract_layers(path):
    img = ImageInput.open(path)
    channels = img.spec().channelnames
    layers = OrderedDict()
    for i in channels:
        layername = ".".join(i.split(".")[:-1])
        for k,v in channel_rename_map.items():
            layername = layername.replace(k,v)
        if layername == "Combined":
            layername = "rgba"
        if layername not in layers:
            layers[layername] = [i]
        else:
            layers[layername].append(i)
    img.close()
    return layers


def convert_to_multipart(path, out_path=None, write=True, replace=False):
    """
    convert a given exr to a multipart exr
    """
    if is_multipart(path):
        return
    buf_spec = OrderedDict()
    layers = extract_layers(path)
    img = ImageBuf(path)
    for layer_name, channels in layers.items():
        new_channels = []
        for ch in channels:
            ext = ch.split(".")[-1]
            new_channels.append(layer_name + "."+ ext)
        buf = oiio.ImageBuf()
        ImageBufAlgo.channels(buf, img, tuple(channels), tuple(new_channels))
        
        bit_depth = oiio.HALF
        
        # check if we want to keeps this si as float
        for float_layer in float_layers:
            if layer_name.lower().__contains__(float_layer):
                bit_depth = oiio.FLOAT
                break
        buf.set_write_format(bit_depth)
        spec = oiio.ImageSpec(img.spec().width, img.spec().height, len(new_channels), bit_depth)
        spec.channelnames = tuple(new_channels)
        # Todo check roi and add it
        
        # fix metadata to make crypos work
        if layer_name.lower().__contains__("crypto") or layer_name.lower().__contains__("rgba"):
            extra_attribs = img.spec().extra_attribs
            for i in range(len(extra_attribs)):
                if fnmatch.filter([extra_attribs[i].name.lower()], "crypto*name"):
                    value = extra_attribs[i].value
                    for k,v in channel_rename_map.items():
                        value = value.replace(k,v)
                    extra_attribs.attribute(extra_attribs[i].name, value)
            spec.extra_attribs = extra_attribs
        buf_spec[buf] = spec
    
    if write:
        if not out_path:
            out_path = construct_out_path(path)
        out = ImageOutput.create(out_path)
        count = 0
        for b,s in buf_spec.items():
            if count>0:
                p = out.open(out_path, s, "AppendSubimage")
            else:
                p = out.open(out_path, tuple(buf_spec.values()))
            count += 1
            b.write(out)
        out.close()
        if replace:
            try:
                os.unlink(path)
                os.rename(out_path, path)
            except Exception as e:
                print(e)
    

convert_to_multipart(path, replace=False)