Source Filmmaker: Script to set light colors using a color picker

set_sfm_light_colors.py

# set_sfm_light_colors.py
#
# Source Filmmaker script.
#
# Finds every animation set in the current document that targets a light
# (ambient / directional / point / spot), shows a checkbox list so you can
# pick which ones to touch, then lets you pick a single RGB color (via a
# native color picker or manual R/G/B fields) and writes that color onto
# every selected light's color_red / color_green / color_blue controls.
#
# Install: drop this file in
#   .../SourceFilmmaker/game/usermod/scripts/sfm/
# (or your mod's scripts/sfm folder) and it will show up under
# Main Menu > Scripts. It can also be pasted into and run from the
# Script Editor directly.
#
# Notes on the SFM data model, since none of this is officially documented:
#   - A light's animation set can be told apart from a model's (or anything
#     else's) by checking animSet.HasAttribute("light").
#   - Light color is NOT a single vs.Color control - it's split into three
#     separate float controls named "color_red", "color_green" and
#     "color_blue" (plus "color_alpha", which this script leaves alone),
#     each living in the animation set's control group and each a
#     0.0-1.0 float, just like every other animatable light property
#     (intensity, radius, etc).
#   - Those float controls are found the same way you'd find any other
#     nested control: animSet.GetRootControlGroup().FindControlByName(name, True)
#   - Because we only want a single, constant color rather than an
#     animated one, every existing key on each control's log is removed
#     and replaced with one key at time zero, which holds for the whole
#     clip.

import sfmApp
import sys
import vs
import re
from PySide import QtGui

COLOR_RED = "color_red"
COLOR_GREEN = "color_green"
COLOR_BLUE = "color_blue"

COLOR_CONTROL_NAMES = (COLOR_RED, COLOR_GREEN, COLOR_BLUE)


def natural_name_key(element):
    """Sort names containing numbers the way people expect (light9, light10)."""
    return [int(part) if part.isdigit() else part
            for part in re.split(r'(\d+)', element.GetName().lower())]


def get_all_animation_sets_direct():
    """Uses the direct sfmApp backend array to find sets regardless of focus."""
    shots = sfmApp.GetShots()
    if not shots:
        return []

    found_sets = []
    for shot in shots:
        if hasattr(shot, 'animationSets'):
            for i in range(len(shot.animationSets)):
                anim_set = shot.animationSets[i]
                if anim_set not in found_sets:
                    found_sets.append(anim_set)
    return found_sets


def get_light_animation_sets():
    """Filters every animation set in the document down to the ones that target a light."""
    lights = []
    for anim_set in get_all_animation_sets_direct():
        try:
            if anim_set.HasAttribute("light"):
                lights.append(anim_set)
        except Exception:
            continue
    lights.sort(key=natural_name_key)
    return lights


def set_light_color(anim_set, r, g, b):
    """Overwrites color_red/green/blue on one light animation set with a single constant keyframe.

    r, g, b are expected to be floats in the 0.0-1.0 range (matching the
    range SFM's own color sliders use). Returns a list of any control
    names that could not be found on this light, so the caller can warn
    about them instead of silently skipping.
    """
    root_group = anim_set.GetRootControlGroup()
    channel_values = {
        COLOR_RED: r,
        COLOR_GREEN: g,
        COLOR_BLUE: b,
    }

    missing = []

    for control_name in COLOR_CONTROL_NAMES:
        control = root_group.FindControlByName(control_name, True)
        if not control:
            missing.append(control_name)
            continue

        layer = control.channel.log.layers[0]

        # Clear every existing key on this control.
        times = vs.tier1.CUtlVectorTime()
        values = vs.tier1.CUtlVectorFloat()
        layer.GetAllKeys(times, values)
        for i in range(times.Count()):
            layer.RemoveKeys(times[i])

        # Write a single key at time zero, which holds constant for the whole clip.
        layer.InsertKey(vs.DmeTime_t(0), channel_values[control_name])

    return missing


def refresh_sfm_view():
    """Force the engine to re-evaluate and render the current frame.

    sfm.GetCurrentFrame()/SetCurrentFrame() are a shot-relative frame count
    used for scripted operations (Move, Rotate, GenerateSamples) - they do
    not drive what the viewport renders. sfmApp.GetHeadTimeInFrames()/
    SetHeadTimeInFrames() are the actual movie-timeline playhead, i.e. what
    changes when you step a frame or drag the scrubber. SFM optimizes out a
    request to set the head to its already-current value, so we nudge the
    head to a neighbouring frame and back, exactly as manually stepping the
    playhead does, to force a real re-render.
    """
    try:
        current_head = sfmApp.GetHeadTimeInFrames()
        neighbor_head = current_head - 1 if current_head > 0 else current_head + 1
        sfmApp.SetHeadTimeInFrames(neighbor_head)
        sfmApp.ProcessEvents()
        sfmApp.SetHeadTimeInFrames(current_head)
        sfmApp.ProcessEvents()
    except Exception:
        pass


def apply_light_color_to_lights(lights, color):
    """Apply ``color`` to ``lights`` and immediately refresh the viewport."""
    r, g, b = color.redF(), color.greenF(), color.blueF()
    changed = 0
    all_missing = {}

    for light in lights:
        try:
            missing = set_light_color(light, r, g, b)
        except Exception as error:
            sys.stderr.write("Failed to update light '%s': %s\n" % (
                light.GetName(), str(error)
            ))
            all_missing[light.GetName()] = ["(error: %s)" % str(error)]
            continue

        if missing:
            all_missing[light.GetName()] = missing
        else:
            changed += 1

    refresh_sfm_view()
    return changed, all_missing


def capture_color_keys(lights):
    """Copy each color control's first log layer for a reversible preview."""
    snapshot = []
    for light in lights:
        controls = {}
        root_group = light.GetRootControlGroup()
        for control_name in COLOR_CONTROL_NAMES:
            try:
                control = root_group.FindControlByName(control_name, True)
                if not control:
                    continue
                layer = control.channel.log.layers[0]
                times = vs.tier1.CUtlVectorTime()
                values = vs.tier1.CUtlVectorFloat()
                layer.GetAllKeys(times, values)
                controls[control_name] = [
                    (times[i].GetTenthsOfMS(), float(values[i]))
                    for i in range(times.Count())
                ]
            except Exception as error:
                sys.stderr.write("Could not snapshot '%s' on '%s': %s\n" % (
                    control_name, light.GetName(), str(error)
                ))
        snapshot.append((light, controls))
    return snapshot


def restore_color_keys(snapshot):
    """Restore a preview snapshot, including the original empty log layers."""
    for light, controls in snapshot:
        root_group = light.GetRootControlGroup()
        for control_name, keys in controls.iteritems():
            try:
                control = root_group.FindControlByName(control_name, True)
                if not control:
                    continue
                layer = control.channel.log.layers[0]
                times = vs.tier1.CUtlVectorTime()
                values = vs.tier1.CUtlVectorFloat()
                layer.GetAllKeys(times, values)
                for i in range(times.Count()):
                    layer.RemoveKeys(times[i])
                for time_in_tenths_ms, value in keys:
                    layer.InsertKey(vs.DmeTime_t(time_in_tenths_ms), value)
            except Exception as error:
                sys.stderr.write("Could not restore '%s' on '%s': %s\n" % (
                    control_name, light.GetName(), str(error)
                ))
    refresh_sfm_view()


def get_light_color(anim_set):
    """Return the RGB value evaluated at the beginning of each control's clip.

    Animated light colors are stored in the channel log.  An unanimated
    channel has no keys, however, and SFM evaluates an empty float log as
    0.0; in that case the control's backing value is the actual light color.
    """
    root_group = anim_set.GetRootControlGroup()
    channel_values = {
        COLOR_RED: 0,
        COLOR_GREEN: 0,
        COLOR_BLUE: 0,
    }
    for control_name in COLOR_CONTROL_NAMES:
        control = root_group.FindControlByName(control_name, True)
        if not control:
            continue

        log = control.channel.GetLog()
        has_keys = False
        for layer_index in range(log.GetNumLayers()):
            times = vs.tier1.CUtlVectorTime()
            values = vs.tier1.CUtlVectorFloat()
            log.GetLayer(layer_index).GetAllKeys(times, values)
            if times.Count() > 0:
                has_keys = True
                break

        if has_keys:
            # Channel-log times are relative to their owner clip. Time zero
            # is therefore the start of that clip, regardless of where it is
            # placed in the shot or timeline.
            value_channel = log.GetValue(vs.DmeTime_t(0))
        else:
            value_channel = control.GetValue("value")
        channel_values[control_name] = max(0, min(255, int(round(255 * value_channel))))
    output = "(R:" + str(channel_values[COLOR_RED]) + ", G:" + str(channel_values[COLOR_GREEN]) + ", B:" + str(channel_values[COLOR_BLUE]) + ")"
    return output


def is_light_enabled(anim_set):
    """Read the Animation Set Editor eye state from its companion DmeDag.

    The projected-light object does not carry this flag.  SFM instead keeps
    an ordinary ``DmeDag`` with the same name for the Animation Set Editor;
    its ``visible`` attribute changes when its eye icon is toggled.
    """
    try:
        light_name = anim_set.GetName()
        maximum = vs.g_pDataModel.GetElementsAllocatedSoFar()
        handle = vs.g_pDataModel.FirstAllocatedElement()
        for _ in range(maximum):
            element = vs.g_pDataModel.GetElement(handle)
            handle = vs.g_pDataModel.NextAllocatedElement(handle)
            if not element:
                continue
            if (element.GetTypeString() == "DmeDag" and
                    element.GetName() == light_name and
                    element.HasAttribute("visible")):
                return bool(element.GetValue("visible"))
    except Exception:
        pass

    # A missing companion is unusual; default to selected rather than hide a
    # light from the user because its editor state could not be read.
    return True


class LightColorDialog(QtGui.QDialog):
    """Checkbox list of light animation sets plus a color picker / RGB fields."""

    def __init__(self, light_sets, parent=None):
        super(LightColorDialog, self).__init__(parent)
        self.setWindowTitle("Set Light Colors")
        self.setMinimumWidth(360)

        self.light_sets = light_sets
        self.checkboxes = []
        self.current_color = QtGui.QColor(255, 255, 255)
        self._original_color_keys = capture_color_keys(light_sets)
        self._has_preview = False
        self.was_cancelled = False

        layout = QtGui.QVBoxLayout(self)

        layout.addWidget(QtGui.QLabel("Lights in this scene:"))

        select_row = QtGui.QHBoxLayout()
        select_all_btn = QtGui.QPushButton("Select All")
        select_none_btn = QtGui.QPushButton("Select None")
        select_all_btn.clicked.connect(self.select_all)
        select_none_btn.clicked.connect(self.select_none)
        select_row.addWidget(select_all_btn)
        select_row.addWidget(select_none_btn)
        select_row.addStretch(1)
        layout.addLayout(select_row)

        list_container = QtGui.QWidget()
        list_layout = QtGui.QVBoxLayout(list_container)
        list_layout.setContentsMargins(4, 4, 4, 4)
        for light in self.light_sets:
            light_color = get_light_color(light)
            light_is_enabled = is_light_enabled(light)
            disabled_suffix = "" if light_is_enabled else " (disabled)"
            light_name = light.GetName() + disabled_suffix + " " + light_color
            checkbox = QtGui.QCheckBox(light_name)
            checkbox.setChecked(light_is_enabled)
            if not light_is_enabled:
                checkbox.setStyleSheet(
                    "QCheckBox { color: rgba(255, 255, 255, 85); }"
                )
            list_layout.addWidget(checkbox)
            self.checkboxes.append((checkbox, light))
        list_layout.addStretch(1)

        scroll = QtGui.QScrollArea()
        scroll.setWidget(list_container)
        scroll.setWidgetResizable(True)
        scroll.setMinimumHeight(200)
        layout.addWidget(scroll)

        layout.addWidget(QtGui.QLabel("Color:"))

        color_row = QtGui.QHBoxLayout()

        self.swatch = QtGui.QLabel()
        self.swatch.setFixedSize(40, 24)
        self.swatch.setFrameShape(QtGui.QFrame.Box)
        color_row.addWidget(self.swatch)

        pick_btn = QtGui.QPushButton("Pick Color...")
        pick_btn.clicked.connect(self.pick_color)
        color_row.addWidget(pick_btn)

        color_row.addSpacing(12)

        color_row.addWidget(QtGui.QLabel("R"))
        self.r_spin = QtGui.QSpinBox()
        self.r_spin.setRange(0, 255)
        self.r_spin.setValue(255)
        color_row.addWidget(self.r_spin)

        color_row.addWidget(QtGui.QLabel("G"))
        self.g_spin = QtGui.QSpinBox()
        self.g_spin.setRange(0, 255)
        self.g_spin.setValue(255)
        color_row.addWidget(self.g_spin)

        color_row.addWidget(QtGui.QLabel("B"))
        self.b_spin = QtGui.QSpinBox()
        self.b_spin.setRange(0, 255)
        self.b_spin.setValue(255)
        color_row.addWidget(self.b_spin)

        for spin in (self.r_spin, self.g_spin, self.b_spin):
            spin.valueChanged.connect(self.spins_changed)

        layout.addLayout(color_row)

        self.update_swatch()

        button_row = QtGui.QHBoxLayout()
        preview_btn = QtGui.QPushButton("Preview")
        preview_btn.clicked.connect(self.preview)
        button_row.addWidget(preview_btn)
        button_row.addStretch(1)

        button_box = QtGui.QDialogButtonBox(QtGui.QDialogButtonBox.Ok | QtGui.QDialogButtonBox.Cancel)
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.cancel)
        button_row.addWidget(button_box)
        layout.addLayout(button_row)

    def select_all(self):
        for checkbox, _ in self.checkboxes:
            checkbox.setChecked(True)

    def select_none(self):
        for checkbox, _ in self.checkboxes:
            checkbox.setChecked(False)

    def update_swatch(self):
        self.swatch.setStyleSheet(
            "background-color: rgb(%d, %d, %d);" % (
                self.current_color.red(),
                self.current_color.green(),
                self.current_color.blue(),
            )
        )

    def spins_changed(self, _value):
        self.current_color = QtGui.QColor(self.r_spin.value(), self.g_spin.value(), self.b_spin.value())
        self.update_swatch()

    def pick_color(self):
        color = QtGui.QColorDialog.getColor(self.current_color, self, "Choose Light Color")
        if not color.isValid():
            return

        self.current_color = color

        for spin in (self.r_spin, self.g_spin, self.b_spin):
            spin.blockSignals(True)
        self.r_spin.setValue(color.red())
        self.g_spin.setValue(color.green())
        self.b_spin.setValue(color.blue())
        for spin in (self.r_spin, self.g_spin, self.b_spin):
            spin.blockSignals(False)

        self.update_swatch()

    def selected_lights(self):
        return [light for checkbox, light in self.checkboxes if checkbox.isChecked()]

    def selected_color(self):
        return self.current_color

    def preview(self):
        """Apply the selected color without closing the dialog."""
        selected = self.selected_lights()
        if not selected:
            return
        self._has_preview = True
        apply_light_color_to_lights(selected, self.selected_color())

    def cancel(self):
        """Discard any preview edits when the Cancel button is used."""
        if self._has_preview:
            restore_color_keys(self._original_color_keys)
        self.was_cancelled = True
        # SFM's old PySide binding errors when Python calls reject() or
        # done(). Closing the modal widget itself works without that binding.
        self.close()


def run_light_color_tool():
    if not sfmApp.HasDocument():
        QtGui.QMessageBox.warning(None, "Light Color Error", "No current document open in SFM!")
        return

    light_sets = get_light_animation_sets()
    if not light_sets:
        QtGui.QMessageBox.warning(
            None,
            "Light Color Error",
            "No light animation sets could be found anywhere in this project file."
        )
        return

    dialog = LightColorDialog(light_sets)
    dialog_result = dialog.exec_()
    if dialog.was_cancelled or dialog_result != QtGui.QDialog.Accepted:
        return

    selected = dialog.selected_lights()
    if not selected:
        QtGui.QMessageBox.information(None, "Light Color", "No lights were selected, nothing was changed.")
        return

    changed, all_missing = apply_light_color_to_lights(selected, dialog.selected_color())

    info_msg = "Updated color on %d of %d selected light(s)." % (changed, len(selected))
    if all_missing:
        info_msg += "\n\nSome lights had issues:\n"
        for name, missing in all_missing.iteritems():
            info_msg += "- %s: %s\n" % (name, ", ".join(missing))

    QtGui.QMessageBox.information(None, "Light Color Update Complete", info_msg)


run_light_color_tool()