blender_3d_ru_ggr_import.py

#!/usr/local/bin/python3
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import os


bl_info = {
    "name": "Load GGR",
    "author": "www.blender-3d.ru",
    "version": (1, 0),
    "blender": (2, 75, 0),
    "location": "Node Editor > Operators",
    "description": "Loads *.ggr file on a node's color ramp",
    "warning": "",
    "wiki_url": "",
    "category": "Node",
}


class GgrColor:
    def __init__(self):
        self.r = 0.0
        self.g = 0.0
        self.b = 0.0
        self.a = 0.0

    def as_tuple(self):
        return (self.r, self.g, self.b, self.a)

    def __eq__(self, other):
        return (self.r == other.r) and (self.g == other.g) and (self.b == other.b) and (self.a == other.a)

    def __str__(self):
        return "Color(%g, %g, %g, %g)" % (self.r, self.g, self.b, self.a)


class GgrSegment:
    def __init__(self):
        self.left_point  = 0.0
        self.mid_point   = 0.0
        self.right_point = 0.0

        self.left_color  = GgrColor()
        self.left_color_type = 0

        self.right_color = GgrColor()
        self.right_color_type = 0

        # Endpoint Color type
        #   0 = "fixed"
        #   1 = "foreground"
        #   2 = "foreground transparent"
        #   3 = "background"
        #   4 = "background transparent"
        #
        self.color_type = 0

        # Blending function
        #   0 = "linear"
        #   1 = "curved"
        #   2 = "sinusoidal"
        #   3 = "spherical (increasing)"
        #   4 = "spherical (decreasing)"
        #
        self.blending_type = 0

        # Coloring type
        #   0 = "RGB"
        #   1 = "HSV CCW"
        #   2 = "HSV CW"
        #
        self.coloring_type = 0

    def length(self):
        return self.right_point - self.left_point

    def calc_mid_point(self):
        return self.left_point + self.length() / 2.0

    def need_mid_point(self):
        return abs(self.calc_mid_point() - self.mid_point) > 0.01


class Ggr:
    def __init__(self):
        self.name = None
        self.segments = []


# File format: https://github.com/mirecta/gimp/blob/master/devel-docs/ggr.txt
#
#  Line 1: "GIMP Gradient"
#  Line 2: "Name: " followed by the name of the gradient
#  Line 3: the number of segments 
#
#  Field      Meaning
#    0          Left endpoint coordinate
#    1          Midpoint coordinate 
#    2          Right endpoint coordinate
#    3          Left endpoint R
#    4          Left endpoint G
#    5          Left endpoint B
#    6          Left endpoint A
#    7          Right endpoint R
#    8          Right endpoint G
#    9          Right endpoint B
#   10          Right endpoint A
#   11          Blending function type
#   12          Coloring type
#   13          Left endpoint color type
#   14          Right endpoint color type
#
def ggr_parse(ggr_file):
    line_1 = ggr_file.readline().strip()
    if line_1 != "GIMP Gradient":
        print("Not a GIMP Gradient file!")
    else:
        ggr = None
        ggr_format = None
        ggr_num_entries = 0
        ggr_name = ggr_file.name

        line_2 = ggr_file.readline().strip()
        if line_2.startswith("Name: "):
            ggr_format = 'NEW'
            ggr_name        = line_2.split(":")[1].strip().replace("_", " ")
            ggr_num_entries = int(ggr_file.readline().strip())
        else:
            ggr_format = 'OLD'
            ggr_num_entries = int(line_2)

        if ggr_format is not None and ggr_num_entries:
            ggr = Ggr()
            ggr.name = ggr_name

            for i in range(ggr_num_entries):
                entry = ggr_file.readline().strip().split()

                ggr_segment = GgrSegment()

                ggr_segment.left_point  = float(entry[0])
                ggr_segment.mid_point   = float(entry[1])
                ggr_segment.right_point = float(entry[2])

                ggr_segment.left_color.r = float(entry[3])
                ggr_segment.left_color.g = float(entry[4])
                ggr_segment.left_color.b = float(entry[5])
                ggr_segment.left_color.a = float(entry[6])

                ggr_segment.right_color.r = float(entry[7])
                ggr_segment.right_color.g = float(entry[8])
                ggr_segment.right_color.b = float(entry[9])
                ggr_segment.right_color.a = float(entry[10])

                if ggr_format == 'NEW':
                    if len(entry) == 13:
                        ggr_segment.blending_type = int(entry[11])
                        ggr_segment.coloring_type = int(entry[12])

                    if len(entry) == 15:
                        ggr_segment.left_color_type = int(entry[13])
                        ggr_segment.right_color_type = int(entry[14])

                ggr.segments.append(ggr_segment)

        if ggr:
            valid = True

            num_markers = 0
            for i,seg in enumerate(ggr.segments):
                add_left_marker = True
                if i:
                    prev_seg = ggr.segments[i-1]
                    if prev_seg.right_color == seg.left_color:
                        add_left_marker = False
                num_markers += 2 if add_left_marker else 1

            if num_markers >= 32:
                valid = False
                print("Can't load gradient: maximum 32 elements!")

            if valid:
                return ggr


def ggr_info(ggr):
    print('Gradient Name: "%s"' % ggr.name)
    print('Segments: %i' % len(ggr.segments))

    for i,seg in enumerate(ggr.segments):
        print("Segment %i:" % i)
        print("  %g - %g%s - %g" % (seg.left_point, seg.mid_point, "*" if seg.need_mid_point() else "", seg.right_point))
        print("  %s - %s" % (seg.left_color, seg.right_color))


def ggr_load_gradients(dirpath):
    import pathlib
    import bpy

    ggrs = bpy.context.scene.blender_3d_ru.ggr.gradients

    for filepath in sorted(pathlib.Path(dirpath).iterdir()):
        if filepath.with_suffix('.ggr'):
            with filepath.open(mode='r') as ggr_file:
                ggr = ggr_parse(ggr_file)
                if not ggr:
                    print("Incorrect GIMP Gradient file: %s" % filepath)
                else:
                    ggr_item = ggrs.add()
                    if ggr.name:
                        ggr_item.name = ggr.name
                    else:
                        ggr_item.name = filepath.name
                    ggr_item.filepath = str(filepath)

                    for i,seg in enumerate(ggr.segments):
                        # Check if we need to add both endpoints
                        add_left_marker = True
                        if i:
                            prev_seg = ggr.segments[i-1]
                            if prev_seg.right_color == seg.left_color:
                                add_left_marker = False

                        if add_left_marker:
                            color_item = ggr_item.colors.add()
                            color_item.color = seg.left_color.as_tuple()

                        color_item = ggr_item.colors.add()
                        color_item.color = seg.right_color.as_tuple()


def ggr_load_file_on_node(filepath, node):
    with open(filepath, 'r') as ggr_file:
        ggr = ggr_parse(ggr_file)
        if ggr:
            ramp = None

            node.label = ggr.name

            if node.bl_idname.startswith('VRayNode'):
                ramp = node.texture.color_ramp
            else:
                ramp = node.color_ramp

            # Remove existing elements
            # NOTE: One element will be left
            for el in reversed(ramp.elements[1:]):
                ramp.elements.remove(el)

            # TODO: HSV / HSL
            ramp.color_mode = 'RGB'

            # TODO: Match mode to the most used in seg modes
            ramp.interpolation = 'LINEAR'

            el_index = 0
            for i,seg in enumerate(ggr.segments):
                right_marker_offset = -0.001

                # Check if we need to add both endpoints
                add_left_marker = True
                if i:
                    prev_seg = ggr.segments[i-1]
                    if prev_seg.right_color == seg.left_color:
                        add_left_marker = False
                        # No need to offset the marker
                        # if we're adding only one
                        right_marker_offset = 0.0

                if add_left_marker:
                    # Reuse element that was left after removal
                    el = ramp.elements[0] if i == 0 else ramp.elements.new(seg.left_point)
                    el.color = seg.left_color.as_tuple()

                el = ramp.elements.new(seg.right_point + right_marker_offset)
                el.color = seg.right_color.as_tuple()

                if seg.need_mid_point():
                    # Evaluate color in the middle of the interval
                    mid_color = ramp.evaluate(seg.calc_mid_point())

                    # Create an element with this value at the
                    # actual midpoint coordinate
                    el = ramp.elements.new(seg.mid_point)
                    el.color = mid_color


def register():
    import bpy

    class GgrOpLoadSelected(bpy.types.Operator):
        bl_idname = "blender_3d_ru.load_ggr_selected"
        bl_label  = "GGR: Load Selected on a Node"

        @classmethod
        def poll(cls, context):
            space = context.space_data
            return space.type == 'NODE_EDITOR'

        def execute(self, context):
            if hasattr(context, 'active_node'):
                node = context.active_node
                if node:
                    ggrs = context.scene.blender_3d_ru.ggr

                    index = ggrs.gradients_selected
                    if index >= 0:
                        ggr_item = ggrs.gradients[index]

                        ggr_load_file_on_node(ggr_item.filepath, node)

            return {'FINISHED'}

    class GgrOpLoad(bpy.types.Operator):
        bl_idname = "blender_3d_ru.load_ggr"
        bl_label  = "GGR: Load from File on a Node"

        filepath = bpy.props.StringProperty(name="File path (*.ggr)", subtype='FILE_PATH')
        context  = None

        @classmethod
        def poll(cls, context):
            space = context.space_data
            return space.type == 'NODE_EDITOR'

        def invoke(self, context, event):
            # NOTE: A bit hacky - we store the original contex,
            # because file browser alters it and we recieve wrong
            # context in execute()
            #
            self.context = context

            context.window_manager.fileselect_add(self)

            return {'RUNNING_MODAL'}

        def execute(self, context):
            ggr_load_file_on_node(self.filepath, self.context.active_node)
            return {'FINISHED'}

    class GgrOpLoadFromFile(bpy.types.Operator):
        bl_idname = "blender_3d_ru.load_ggr_from_file"
        bl_label  = "GGR: Load from File on a Node"
        bl_description = "Load gradient on a selected node"

        filepath = bpy.props.StringProperty(name="File path (*.ggr)", subtype='FILE_PATH')

        def execute(self, context):
            active_node = None

            # Find selected node
            for window in context.window_manager.windows:
                for area in window.screen.areas:
                    for space in area.spaces:
                        if space.type == 'NODE_EDITOR':
                            ntree = space.node_tree
                            if ntree and ntree.nodes:
                                # Active node is always last
                                selected_node = ntree.nodes[-1]

                                # Use only if node has color ramp
                                accept = False
                                if selected_node.bl_idname.startswith('VRayNode'):
                                    if hasattr(selected_node, 'texture'):
                                        accept = True
                                elif hasattr(selected_node, 'color_ramp'):
                                    accept = True

                                if accept:
                                    active_node = selected_node
                                    break

            if active_node:
                ggr_load_file_on_node(self.filepath, active_node)

            return {'FINISHED'}

    class GgrGradItem(bpy.types.PropertyGroup):
        color = bpy.props.FloatVectorProperty(
            min     = 0.0,
            max     = 1.0,
            subtype = 'COLOR',
            size    = 4
        )

    class GgrGradient(bpy.types.PropertyGroup):
        name     = bpy.props.StringProperty()
        filepath = bpy.props.StringProperty()
        colors   = bpy.props.CollectionProperty(type=GgrGradItem)

    class GgrGradients(bpy.types.PropertyGroup):
        gradients = bpy.props.CollectionProperty(type=GgrGradient)
        gradients_selected = bpy.props.IntProperty(
            name    = "Scene Index",
            default = -1,
            min     = -1,
        )

    class Blender3dRu(bpy.types.PropertyGroup):
        ggr = bpy.props.PointerProperty(type=GgrGradients)

    class GgrListItem(bpy.types.UIList):
        def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
            split = layout.split(percentage=0.35)

            col = split.column()
            col.label(item.name)

            col = split.column()
            sub = col.row()

            row = sub.row(align=True)
            for col_item in item.colors:
                 row.prop(col_item, 'color', text="")

            row = sub.row(align=True)
            op = row.operator('blender_3d_ru.load_ggr_from_file', text="", icon='COLOR')
            op.filepath = item.filepath

    class GgrOpLoadDirectory(bpy.types.Operator):
        bl_idname = "blender_3d_ru.load_ggr_directory"
        bl_label  = "GGR: Load Directory"

        directory = bpy.props.StringProperty(name="Directory path", subtype='DIR_PATH')

        def invoke(self, context, event):
            context.window_manager.fileselect_add(self)
            return {'RUNNING_MODAL'}

        def execute(self, context):
            gradients = context.scene.blender_3d_ru.ggr.gradients

            # Clear existing
            for i in reversed(range(len(gradients))):
                gradients.remove(i)

            # Load new
            ggr_load_gradients(self.directory)

            return {'FINISHED'}

    class GgrPanel(bpy.types.Panel):
        bl_space_type  = 'PROPERTIES'
        bl_region_type = 'WINDOW'
        bl_context     = 'scene'
        bl_label       = "Gradients"

        def draw(self, context):
            ggr = context.scene.blender_3d_ru.ggr

            self.layout.operator('blender_3d_ru.load_ggr_directory', icon='FILE_FOLDER')
            self.layout.template_list("GgrListItem", "", ggr, 'gradients', ggr, 'gradients_selected', rows= 3)

    _classes = (
        GgrGradItem,
        GgrGradient,
        GgrGradients,
        GgrOpLoad,
        GgrOpLoadDirectory,
        GgrOpLoadSelected,
        GgrOpLoadFromFile,
        GgrListItem,
        GgrPanel,
        Blender3dRu,
    )

    for _cls in _classes:
        bpy.utils.register_class(_cls)

    bpy.types.Scene.blender_3d_ru = bpy.props.PointerProperty(name="Blender 3D RU", type=Blender3dRu)


def unregister():
    pass


def ggr_test():
    ggr = Ggr()
    ggr_seg = GgrSegment()
    ggr_seg.mid_point = 0.5
    ggr_seg.right_point = 1.0
    ggr.segments.append(ggr_seg)
    ggr_info(ggr)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('ggr_file', nargs='?', type=argparse.FileType('r'))
    parser.add_argument('--test', action='store_true', default=False)

    args = parser.parse_args()
    if not args.ggr_file:
        parser.print_help()
    else:
        if args.test:
            ggr_test()

        ggr = ggr_parse(args.ggr_file)
        ggr_info(ggr)