hsd2ora.py

import os
import tempfile
import json
import pyora
import io
import shutil
import math
import fpng_py
import argparse
import sys
import uuid
from pathlib import Path
from batch_processing import Batch
from PIL import Image
from czipfile import ZipFile #https://github.com/ziyuang/czipfile
import xml.etree.ElementTree

#for some reason, this tries to write boolean values to the xml. so i've modified it to convert them to strings first.
def override_escape_attrib(text):
    # escape attribute value
    try:
        if (isinstance(text, bool)):
            text = str(text).lower()
        if "&" in text:
            text = text.replace("&", "&")
        if "<" in text:
            text = text.replace("<", "&lt;")
        if ">" in text:
            text = text.replace(">", "&gt;")
        if "\"" in text:
            text = text.replace("\"", "&quot;")
        # Although section 2.11 of the XML specification states that CR or
        # CR LN should be replaced with just LN, it applies only to EOLNs
        # which take part of organizing file into lines. Within attributes,
        # we are replacing these with entity numbers, so they do not count.
        # http://www.w3.org/TR/REC-xml/#sec-line-ends
        # The current solution, contained in following six lines, was
        # discussed in issue 17582 and 39011.
        if "\r" in text:
            text = text.replace("\r", "&#13;")
        if "\n" in text:
            text = text.replace("\n", "&#10;")
        if "\t" in text:
            text = text.replace("\t", "&#09;")
        return text
    except (TypeError, AttributeError):
        _raise_serialization_error(text)

xml.etree.ElementTree._escape_attrib = override_escape_attrib

#makes pyora use czipfile to save, so it goes faster 
def overridesave(self, path_or_file, composite_image=None, use_original=False):
    """Save the current project state to an ORA file.

    Args:
        path (str): path to the ora file to save
        composite_image (PIL.Image()): PIL Image() object of the composite rendered canvas. It is used to
            create the mergedimage full rendered preview, as well as the thumbnail image. If not provided,
            one will be generated by pyora's Render() class by stacking all of the layers in the project.
            Note that the image you pass may be modified during this process, so if you need to use it elsewhere
            in your code, you should copy() first.
        use_original (bool): If true, and If there was a stored 'mergedimage' already in the file which was opened,
            use that for the 'mergedimage' in the new file, instead of rendering a new one.
    """
    import czipfile
    with czipfile.ZipFile(path_or_file, "w") as zipref:

        zipref.writestr("mimetype", "image/openraster".encode())

       
        if not composite_image:
            if use_original and self._extracted_merged_image:
                composite_image = self._extracted_merged_image
            else:
                # render using our built in library
                r = pyora.Render.Renderer(self)
                composite_image = r.render()
        self._zip_store_image(zipref, "mergedimage.png", composite_image)

        pyora.Render.make_thumbnail(composite_image)  # works in place
        self._zip_store_image(zipref, "Thumbnails/thumbnail.png", composite_image)
        
        
        filename_counter = 0
        for layer in self.children_recursive:
            if layer.type == pyora.TYPE_LAYER:
                new_filename = f"/data/layer{filename_counter}.png"
                layer._elem.attrib["src"] = new_filename
                filename_counter += 1
                #we can now handle using a path to a png file as image data. this negates having to store the whole uncompressed image in memory.
                if isinstance((layer.get_image_data(raw=True)), str):
                    self._zip_store_image(
                        zipref, layer["src"], (Image.open((layer.get_image_data(raw=True))))
                    )
                else:
                    self._zip_store_image(
                        zipref, layer["src"], layer.get_image_data(raw=True)
                    )
                
        zipref.writestr("stack.xml", xml.etree.ElementTree.tostring(self._elem_root, method="xml"))

#fixes a crash when getting z_index
@property
def _override_z_index(self):
    """Get the current z_index

    Get the stacking position of the layer, relative to the group it is in (or the root group).
    Higher numbers are 'on top' of lower numbers. The lowest value is 1.

    Returns:
        int: the z_index of the layer
    """
    if self.parent is None:
        return 1
    return list(reversed(list(self.parent._elem))).index(self._elem) + 1

pyora.Layer.z_index = _override_z_index
pyora.Project.save = overridesave

#all blending modes i was able to select in hipaint are listed. gaps present are as hipaint has provided them.
blendmodes = {  0 : "svg:src-over", 
                1 : "svg:overlay",
                2 : "svg:darken",
                3 : "svg:multiply",
                4 : "svg:color-burn",
                5 : "krita:linear_burn", #krita specific
                6 : "krita:darker color", #krita specific
                7 : "svg:lighten",
                8 : "svg:screen",
                9 : "svg:color-dodge",
                10 : "svg:plus",
                11 : "krita:lighter color", #krita specific
                12 : "svg:soft-light",
                13 : "",
                14 : "svg:hard-light",
                15 : "krita:vivid_light", #krita specific
                16 : "krita:linear light", #krita specific
                17 : "krita:pin_light", #krita specific
                18 : "krita:hard mix", #krita specific
                19 : "svg:difference",
                20 : "krita:exclusion", #krita specific
                21 : "",
                22 : "krita:divide", #krita specific
                23 : "",
                24 : "",
                25 : "svg:hue",
                26 : "svg:saturation",
                27 : "svg:color",
                28 : "svg:luminosity",
                29 : "krita:subtract", #krita specific
                30 : "krita:dissolve", #krita specific
                31 : "",
                32 : "",
                #33 : "group penetrate" #this one is handled differently, as it is not a layer mode in ora, but a distinct attribute
}

def int_to_hex_color(value):
    # Convert signed int to unsigned int
    unsigned_value = value & 0xFFFFFFFF
    
    # Extract the color components
    alpha = (unsigned_value >> 24) & 0xFF
    red = (unsigned_value >> 16) & 0xFF
    green = (unsigned_value >> 8) & 0xFF
    blue = unsigned_value & 0xFF
    
    # Format the hex color code without the alpha
    hex_color = f'#{red:02X}{green:02X}{blue:02X}'
    return hex_color

#takes an hsd file, converts to ora through a bunch of other functions, saves to output 
def convertToORA(filename, output):
    print("Attempting to convert " + str( os.path.abspath(filename)))
    #generate the temp dir we'll be using throughout the function, extract hsd into it + load project json
    tempdir = tempfile.TemporaryDirectory()
    projectDetails = extractProject(filename, tempdir)
    projpath = os.path.join(tempdir.name, "temp")
    
    #create new ora with height taken from hsd json 
    #TO-DO: cropping? i haven't accounted for the vars for it in the json, but it doesn't seem to affect the final image in any of my hsd files
    oraProject = pyora.Project.new(projectDetails['bounds']['canvas-width'], projectDetails['bounds']['canvas-height'])
    
    #background is handled separately from layers in the hsd, so we do it before
    #converting the colour to hex might be unnecessary. test
    bgColor = Image.new('RGB', (projectDetails['bounds']['canvas-width'], projectDetails['bounds']['canvas-height']), int_to_hex_color(projectDetails['background']['bg-color']))
    oraProject.add_layer(bgColor, '/background')
    
    #loop through layers and add them to the project
    for x in projectDetails['layers']:
        generateLayer(x, projectDetails, oraProject, projpath, tempdir)
    
    
        
    #assign layers to their groups, now that all layers and groups are established.
    #pyora is supposed to be able to handle groups automatically by setting the paths correctly, but i couldn't get it to work, so we loop through the layers again and assign children to parents (because i fear what will happen if i try to assign a layer as a child of a layer that doesn't exist yet)
    for x in projectDetails['layers']:
        assignParent(x['filename-id'], x['parent-id'], oraProject)
        
    #now that all layers are in, loop through again. determine clipping sub groups. group them. assign the group the id of their base layer.
    groupClipping(projectDetails, oraProject)

    #TO-DO: take "selected-layer" value from json and set selected layer in the ora  (selected-layer will be set to the uuid of the relevant ora layer). bg layer may require special accommodations, i haven't tested it.
    
    #that should be everything finished. export as osd file. pull composite img directly from hsd.
    oraProject.save(output,composite_image=(Image.open(os.path.join(projpath, "preview"))))
    print("Saved file to " + str( os.path.abspath(output)))


def groupClipping(projectDetails, oraProject):
    innerCounter = 0
    #look through layers. skip layers already assigned to clipping groups, and non-clipping layers.
    for x in projectDetails['layers']:
        #print("scanning layer " + str(x['filename-id']))
        if innerCounter > 0:
            #print("skipping layer")
            innerCounter = innerCounter - 1
            continue
        if x['clip'] != True:
            #print("not a clipping layer. "+str(x['filename-id']))
            continue
        #got a clipping layer. take note of its parent so we know where to assign the clipping group layer later on. also generate the clipping group in ora
        hipaintparent = str(x['parent-id'])
        parentORALayer = oraProject.get_by_uuid(str(x['filename-id'])).parent
        clippingGroup = oraProject.add_group(path="/")
        clippingGroup.name = (projectDetails['layers'][(x['id']-1)]['name']) + (" clipgrouphandler")
        #start tracking the layers we'll group. start with this layer and the layer right below it (the base layer).
        layersToMove = []
        layersToMove.append(projectDetails['layers'][(x['id']-1)])
        layersToMove.append(projectDetails['layers'][(x['id'])])
        #check if next one up is same parent and clipping.
        innerCounter = 1
        while True:
            if (x['id']+innerCounter) > len(projectDetails['layers']):
                #print("prospect is over top layer.")
                innerCounter  = innerCounter - 1
                break
            #print("checking prospect w/ order id: " + str((x['id']+innerCounter)))
            prospectLayer = projectDetails['layers'][(x['id']+innerCounter)]
            if str(prospectLayer['parent-id']) != hipaintparent:
                #print("prospect is not from same parent. "+ str(prospectLayer['filename-id']))
                #print(str(prospectLayer['parent-id']) + ", compared to our golden: "+str(hipaintparent))
                innerCounter  = innerCounter - 1
                break
            if prospectLayer['clip'] == False:
                #print("prospect is not a clipping layer. " + str(prospectLayer['filename-id']))
                innerCounter  = innerCounter - 1
                break
            innerCounter = innerCounter + 1
            layersToMove.append(prospectLayer)
            continue
        #now we know all the layers we want in our group. turrah turrah! note down base layer uuid and z_index.
        baseLayer = oraProject.get_by_uuid(str(projectDetails['layers'][x['id']-1]['filename-id']))
        baseZ = baseLayer.z_index
        baseUUID = baseLayer.uuid
        #then move all layers into clipping group.
        #print("moving these layers into clip group: "+str(layersToMove))
        for m in layersToMove:
            oraProject.move(str(m['filename-id']), clippingGroup.uuid,dst_z_index='above') #dst_z_index?
        #change base layer uuid to random. then assign old base layer uuid to our clipping group, and move it into the old z index.
        oraProject.get_by_uuid(hipaintparent).uuid = str(uuid.uuid4())
        clippingGroup.uuid = baseUUID
        oraProject.move(baseUUID, parentORALayer.uuid, dst_z_index=baseZ)

#assigns a layer in an ora file as a child to another layer using UUIDs
def assignParent(childUUID, parentUUID, oraProject):
    #if we're given an id of 0 or less, the parent layer does not exist, as ensured by extractProject. so we leave it be, as it ought to be in root dir already
    if(parentUUID <= 0):
        return True

    #convert to string b/c they're probably ints from the hsd file
    oraProject.move(str(childUUID), str(parentUUID), dst_z_index='above')
    return True

#takes a layer dict from an hsd json file (converted to a dict beforehand), turns into an ora layer in provided ora project file
#params: the dict of the specific layer, the hsd dict as a whole, the ora project to add the layer to, the path where hsd project files are stored, path to temp dir where layer image files are stored
def generateLayer(layer, hsdDict, oraProject, projpath, layerImgDir): 
    #as far as i can tell, bean-type determines whether it's a paint layer or group
    if layer['bean-type'] == 1:
        new_layer = oraProject.add_group(path="/")
    else:
        #if it's an image layer, get the path of where image data should be
        layerImageFilePath = os.path.join(projpath, ("flayer_"+str(layer['filename-id'])))
        #layer may be empty. check if file id exists. if it does, add it as the image data
        if(os.path.isfile(layerImageFilePath)):
            layerImageDataArray = readLayer(layerImageFilePath, hsdDict, layerImgDir) #returns path to png image data, x offset, y offset 
            invertedYOffset = layerImageDataArray[2]
            new_layer  = oraProject.add_layer(layerImageDataArray[0], "/", offsets=(layerImageDataArray[1], invertedYOffset))
        #or else, if the layer is empty, generate empty image data (so that it doesn't lock up for lack of data to render
        else:
            new_layer  = oraProject.add_layer(Image.new("RGB", (1, 1)), "/", offsets=(0, 0))
    #filename-id is the important id. the regular id value, as far as i can tell, just indicates the order of the layers--which we automatically derive from the order of the entries in the json file, which just so happens to be the same order as the id values
    new_layer.uuid = str(layer['filename-id'])
    new_layer.name = str(layer['name'])
    #now add any attributes
    if layer['visible'] is False:           
        new_layer.visible = False
    new_layer.opacity = layer['opacity']
    #"lock-opacity" in the hsd is whether transparent pixels are locked. i cannot find an attribute in the osd specs that implements this, and since krita does not retain this setting upon saving, i do not think there is one. however, i've put it here anyway, in case it's implemented later. thankfully, it does not affect the final image.
    new_layer["lock-opacity"] = layer['lock-opacity']
    if layer['blend'] == 33:
        #penetrate blend mode is treated as an attribute in ORA
        new_layer.isolated = False;
    else:
        #blend determines blend mode. this is a pain, because they're assigned numbers instead of string based names. i've made a dict with all the numbers assigned their appropriate blending mode name, for ease of use.
        #if you crash because of an unspecified blend mode, add it to the dict above, and then let me know about it so i can do the same
        new_layer["composite-op"] = blendmodes[layer['blend']]
    new_layer["alpha-preserve"] = layer['clip']
    new_layer["edit-locked"] = layer['lock-layer']
    #"is-foreground" ?
    #"is-background" ?
    return True

#reads an hsd layer file and returns a list containing a path to the layer img as png, plus offsets
def readLayer(filename,hsdDict,layerImgDir):
    #the layer file is a zip containing a bitmap. first, we need to extract the zip.
    tempdir = tempfile.TemporaryDirectory()
    with ZipFile(filename) as zf:
        zf.extractall(path=tempdir.name)
    #the file inside will be called 'zip'
    with open((os.path.join(tempdir.name, "zip")), mode="rb") as f:
        buffer_size = 2**10*8
        #we open an 8kb chunk of our file
        chunk = f.read(buffer_size)
        #read header data
        width = int.from_bytes(chunk[8:10], "little")
        height = int.from_bytes(chunk[10:12], "little")
        x_offset = int.from_bytes(chunk[4:6], "little")
        y_offset =int.from_bytes(chunk[6:8], "little")

        #ridiculously high/low values means that there are no header values and it is reading pixel data. how do we avoid this? divide amount of bytes in the layer file by 4. compare that number to the width*height. if they are the same, there is no header data. treat the layer as full canvas, and set whole file to imageData.
        osstat = os.stat((os.path.join(tempdir.name, "zip"))).st_size
        if (hsdDict['bounds']['canvas-width']*hsdDict['bounds']['canvas-height']) <= (osstat/4):
            width = hsdDict['bounds']['canvas-width']
            height = hsdDict['bounds']['canvas-height']
            x_offset = 0
            y_offset = 0
            imageData = bytearray(chunk)
        else:
            #read the rest of the non-header bytes in the chunk. store in bytearray
            imageData = bytearray(chunk[12:(buffer_size + 12)])
        while chunk:
            chunk = f.read(buffer_size)
            imageData.extend(chunk)

        #first, calculate how much image data we actually have. 
        actualRows = (len(imageData)/4)/width
        #we'll probably get a decimal. we want to round up, which necessitates padding with transparent pixels to fill out the last row. we won't bother calculating the exact amount, we'll allow some overshoot and crop it out.
        for _ in range(width):
            imageData.extend((b'\x00\x00\x00\x00')) 
        im = Image.new("RGBA", (width, math.ceil(actualRows)))
        im.frombytes(imageData)
        del imageData
        im = im.transpose(method=Image.FLIP_TOP_BOTTOM)

        #then create a canvas of appropriate size and paste it into that.
        fullDimensions = Image.new("RGBA", (width, height), (0, 0, 0, 0))
        y = height - math.ceil(actualRows)
        area = (0 ,y) #determines upper left pixel position.
        fullDimensions.paste(im, area)
        del im
        #we're hitting memory errors doing all layers in memory. so we save image and pull it back up when needed.
        fpng_py.fpng_encode_image_to_file((os.path.join(layerImgDir.name, (filename+".png"))),fullDimensions.tobytes(),width,height,4)
        y_offset = hsdDict['bounds']['canvas-height'] - y_offset - fullDimensions.height
        del fullDimensions
        productFilename = (os.path.join(layerImgDir.name, (filename+".png")))
        return([productFilename, x_offset, y_offset])

#extracts an hsd project from a provided path, extracts to provided directory, fixes parent ids in project json
#returns the fixed project json as a dict
def extractProject(filename, directory):
    with ZipFile(filename, 'r') as zf:
        zf.extractall(path=directory.name,pwd=b'huion2018')
        #all files will be in a subdir "temp" due to the way hipaint structures project files
    with open(os.path.join(directory.name, "temp", "project.json"), encoding='utf-8') as f:
        projectDetails = json.load(f)
    #hsd files have weird ways of designating that a layer does not have a parent. to get around that, we get a list of every id that actually exists, and run through the list of layers, checking parent IDs against it. if the parent id does not actually exist, we set it to a negative value, so we can later easily check if a layer is on the root layer by seeing if parent-id <= 0
    extantIDs = []
    for x in projectDetails['layers']:
        extantIDs.append(x['filename-id'])
    for x in projectDetails['layers']:
        if x['parent-id'] not in extantIDs:
            x['parent-id'] = -999
    return projectDetails

#batch processing stuff that i threw in and didn't test thoroughly, sorry if there's issues
def parse_command_line(batch: Batch):
    parser = argparse.ArgumentParser(description='Batch input',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Setup custom script arguments here

    return batch.parse_args(parser)


def parallel_process(input: Path, output: Path, args: argparse.Namespace):

    # Load an input file, apply some processing, and save to output file here
    convertToORA(input, output)
    return   


def main(argv):
    
    # Instantiate batch
    batch = Batch(argv)
    batch.set_io_description(input_help='input files', output_help='output directory')
    
    # Parse arguments
    args = parse_command_line(batch)
    
    # Start processing and retrieve the results at the end if any
    data = batch.run(parallel_process, output_ext=".ora")
    
    # clean up residue
    dirpath = Path(os.path.join(args.output, "_batch"))
    if dirpath.exists() and dirpath.is_dir():
        shutil.rmtree(dirpath)


if __name__ == "__main__":
    main(sys.argv[1:])