create-mekanism-modular/tools/ARMOUR_GUI.py

567 lines
22 KiB
Python

#!/usr/bin/env python3
"""Draw geometry onto a body part, unwrap it, and see it on the wearer.
The placement editor moves a module that already exists. This makes one. A gem
that has to follow an arm has to be drawn *by* the arm - MIAPI picks which body
part a model renders under by comparing the model's own `origin` against the
part it is drawing, and it renders under that part's animated pose, so a socket
cut into `left_arm` swings with the arm and one placed in `body` does not.
Armory's gemstone declares no origin at all, which is why it cannot be made to
follow a limb from the outside, and why the socket has to be ours.
So this edits a model file of our own: boxes in the limb's own coordinates,
shown against Armory's plate so they can be lined up with it, unwrapped onto a
texture that is written out beside them.
tools/ARMOUR_GUI.py --jar <armory.jar>
Left is the box list and the part it belongs to, right is the box being edited,
middle is the wearer. Everything is in model pixels, the units the JSON is
written in.
"""
from __future__ import annotations
import argparse
import json
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import guiplatform
guiplatform.configure(prefer=os.environ.get('ARMOUR_GUI_PLATFORM'))
import numpy as np # noqa: E402
from PySide6 import QtCore, QtWidgets # noqa: E402
import pyvistaqt # noqa: E402
import mcmodel as mc # noqa: E402
import ARMOUR_EDITOR as ae # noqa: E402
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ASSETS = os.path.join(REPO, 'src/main/resources/assets/cmmodular')
# Which Armory slot dresses each body part, and the plate to draw as context.
# The key is the `origin` MIAPI matches against, which is the whole point of
# choosing a part: it decides what the geometry is pinned to when the wearer
# moves.
PARTS = {
'left_arm': {'slot': 'arm_left', 'piece': 'chestplate', 'plate': 'arm_left/heavy'},
'right_arm': {'slot': 'arm_right', 'piece': 'chestplate', 'plate': 'arm_right/heavy'},
'left_leg': {'slot': 'leg_left', 'piece': 'pants', 'plate': 'leg_left/heavy'},
'right_leg': {'slot': 'leg_right', 'piece': 'pants', 'plate': 'leg_right/heavy'},
'body': {'slot': 'chest_front', 'piece': 'chestplate', 'plate': 'chest_front/heavy'},
'head': {'slot': 'hat', 'piece': 'helmet', 'plate': 'helmet/heavy'},
# Not a body part at all, but the same choice: `item` is the pass that draws
# the inventory icon, and geometry filed under it is what the icon shows.
# It is on this list because it is the other half of the same decision -
# a socket needs one model on the limb and one on the icon, and they are
# different files with different origins.
'item': {'slot': 'icon', 'piece': None, 'plate': None},
}
WEARER = '(wearer)'
PART_LABELS = {
'left_arm': 'left arm - pauldron', 'right_arm': 'right arm - pauldron',
'left_leg': 'left leg - knee', 'right_leg': 'right leg - knee',
'body': 'chest', 'head': 'helmet', 'item': 'inventory icon',
WEARER: 'the wearer',
}
PIECE_FILES = {
'chestplate': 'data/tm_armory/miapi/modules/armor/chestplate.json',
'pants': 'data/tm_armory/miapi/modules/armor/pants.json',
'helmet': 'data/tm_armory/miapi/modules/armor/helmet.json',
}
# ------------------------------------------------------------------- workpiece
class Workpiece:
"""One cmmodular model file: boxes in a body part's own coordinates."""
def __init__(self, res, part):
self.res, self.part = res, part
self.slot = PARTS[part]['slot']
self.path = os.path.join(
ASSETS, f'models/item/armor/model/{self.slot}/socket/default.json')
self.texture_ref = f'cmmodular:equipment/{self.slot}_socket'
self.texture_path = os.path.join(
ASSETS, f'textures/equipment/{self.slot}_socket.png')
self.doc = self._load()
def _load(self):
if os.path.isfile(self.path):
with open(self.path) as fh:
return json.load(fh)
return {
'comment': f'Socket geometry drawn under {self.part}, so it follows '
f'the limb rather than the torso.',
'texture_size': [16, 16],
'textures': {'0': self.texture_ref, 'particle': self.texture_ref},
'elements': [],
}
@property
def elements(self):
return self.doc.setdefault('elements', [])
def add_box(self, name='socket', lo=(-1.0, -1.0, -1.0), hi=(1.0, 1.0, 1.0)):
self.elements.append({'name': name, 'from': list(lo), 'to': list(hi),
'faces': {}})
self.unwrap()
return len(self.elements) - 1
def remove(self, index):
if 0 <= index < len(self.elements):
self.elements.pop(index)
self.unwrap()
def unwrap(self):
"""Re-cut the texture so every face has somewhere of its own to live."""
if not self.elements:
self.doc['texture_size'] = [16, 16]
return None
size, nets = mc.unwrap(self.elements, texture='#0')
self.doc['texture_size'] = [int(size[0]), int(size[1])]
return size, nets
def save(self, write_template=True):
packed = self.unwrap()
os.makedirs(os.path.dirname(self.path), exist_ok=True)
with open(self.path, 'w') as fh:
fh.write(ae._dumps(self.doc))
fh.write('\n')
written = [self.path]
# Only ever write a template over a texture that is not there yet -
# the guide is scaffolding, and overwriting art someone has painted
# because the box list changed would be the tool destroying the work
# it exists to support.
if write_template and packed and not os.path.isfile(self.texture_path):
size, nets = packed
os.makedirs(os.path.dirname(self.texture_path), exist_ok=True)
mc.unwrap_template(size, nets, self.elements).save(self.texture_path)
written.append(self.texture_path)
return written
def quads(self, res):
model = {'textures': self.doc.get('textures', {}),
'elements': self.elements,
'texture_size': self.doc.get('texture_size')}
return mc.model_quads(model, res)
# ----------------------------------------------------------------------- window
class ArmourGui(QtWidgets.QMainWindow):
def __init__(self, res, args):
super().__init__()
self.res, self.args = res, args
self.setWindowTitle('armour geometry')
self.work = Workpiece(res, args.part)
self.actors = {}
splitter = QtWidgets.QSplitter()
splitter.addWidget(self._left_panel())
self.view = pyvistaqt.QtInteractor(self, rw=guiplatform.render_window())
splitter.addWidget(self.view)
splitter.addWidget(self._right_panel())
splitter.setSizes([230, 900, 250])
self.setCentralWidget(splitter)
self.statusBar().showMessage(f'{self.work.path}')
self.view.set_background(args.background)
self._add_lights()
self._realised = False
def showEvent(self, event):
"""First draw waits for the window.
A QOpenGLWidget has no GL context until it is on screen, and VTK asked
to render before that goes looking for one of its own - which on a
Wayland session means a GLX context that cannot be made current.
"""
super().showEvent(event)
if not self._realised:
self._realised = True
QtCore.QTimer.singleShot(0, self._first_draw)
def _first_draw(self):
# Depth peeling probes the GL context, so it has to wait for one too.
self.view.enable_depth_peeling(number_of_peels=8, occlusion_ratio=0.0)
self.refresh(reset=True)
# ------------------------------------------------------------- panels
def _left_panel(self):
box = QtWidgets.QWidget()
lay = QtWidgets.QVBoxLayout(box)
lay.addWidget(QtWidgets.QLabel('armour parts'))
# One list doing both jobs: the tick says whether a part is drawn, the
# selection says which one the boxes below belong to. They are the same
# question asked twice otherwise - you cannot line a socket up against a
# pauldron you have hidden, and the part you are editing is the one you
# always want on screen, so selecting a row ticks it.
self.parts_list = QtWidgets.QListWidget()
self.parts_list.setFixedHeight(150)
for part in list(PARTS) + [WEARER]:
item = QtWidgets.QListWidgetItem(PART_LABELS.get(part, part))
item.setData(QtCore.Qt.UserRole, part)
item.setFlags(item.flags() | QtCore.Qt.ItemIsUserCheckable)
visible = part == self.args.part or (part == WEARER and self.args.body)
item.setCheckState(QtCore.Qt.Checked if visible else QtCore.Qt.Unchecked)
self.parts_list.addItem(item)
self.parts_list.itemChanged.connect(lambda _: self.refresh())
self.parts_list.currentItemChanged.connect(self._parts_selected)
self.parts_list.setCurrentRow(list(PARTS).index(self.args.part))
self._show_edited_part()
lay.addWidget(self.parts_list)
row = QtWidgets.QHBoxLayout()
for label, state in (('All', QtCore.Qt.Checked), ('None', QtCore.Qt.Unchecked)):
button = QtWidgets.QPushButton(label)
button.clicked.connect(lambda _=None, st=state: self._set_all(st))
row.addWidget(button)
lay.addLayout(row)
note = QtWidgets.QLabel('geometry follows the selected part when the '
'wearer moves')
note.setWordWrap(True)
note.setStyleSheet('color: #888;')
lay.addWidget(note)
lay.addWidget(QtWidgets.QLabel('boxes'))
self.list = QtWidgets.QListWidget()
self.list.currentRowChanged.connect(lambda _: self.refresh())
lay.addWidget(self.list, 1)
for label, slot in (('Add box', self.on_add),
('Duplicate', self.on_duplicate),
('Remove', self.on_remove),
('Unwrap UVs', self.on_unwrap),
('Save model + template', self.on_save)):
button = QtWidgets.QPushButton(label)
button.clicked.connect(slot)
lay.addWidget(button)
self.textured = QtWidgets.QCheckBox('textures')
self.textured.setChecked(not self.args.no_textures)
self.textured.toggled.connect(lambda _: self.refresh())
lay.addWidget(self.textured)
return box
# ------------------------------------------------------- visible parts
def _rows(self):
for i in range(self.parts_list.count()):
yield self.parts_list.item(i)
def _visible(self, part):
for item in self._rows():
if item.data(QtCore.Qt.UserRole) == part:
return item.checkState() == QtCore.Qt.Checked
return False
def _set_all(self, state):
self.parts_list.blockSignals(True)
for item in self._rows():
item.setCheckState(state)
self.parts_list.blockSignals(False)
self._show_edited_part()
self.refresh()
def _show_edited_part(self):
"""The part being edited is never hidden - that would hide the work."""
self.parts_list.blockSignals(True)
for item in self._rows():
part = item.data(QtCore.Qt.UserRole)
font = item.font()
font.setBold(part == self.args.part)
item.setFont(font)
if part == self.args.part:
item.setCheckState(QtCore.Qt.Checked)
self.parts_list.blockSignals(False)
def _parts_selected(self, item, _previous=None):
if item is None:
return
part = item.data(QtCore.Qt.UserRole)
if part == WEARER or part == self.args.part:
self._show_edited_part()
return
self._switch_part(part)
def _right_panel(self):
box = QtWidgets.QWidget()
lay = QtWidgets.QFormLayout(box)
self.name_edit = QtWidgets.QLineEdit()
self.name_edit.editingFinished.connect(self.on_name)
lay.addRow('name', self.name_edit)
self.spins = {}
for key in ('from', 'to'):
for i, axis in enumerate('xyz'):
spin = QtWidgets.QDoubleSpinBox()
spin.setRange(-64.0, 64.0)
spin.setSingleStep(0.25)
spin.setDecimals(3)
spin.valueChanged.connect(self.on_spin)
self.spins[(key, i)] = spin
lay.addRow(f'{key} {axis}', spin)
self.size_label = QtWidgets.QLabel('-')
lay.addRow('size', self.size_label)
self.atlas_label = QtWidgets.QLabel('-')
lay.addRow('texture', self.atlas_label)
return box
def _add_lights(self):
import pyvista as pv
for direction, intensity in (((0.2, -1.0, -0.7), 0.62),
((-0.2, -1.0, 0.7), 0.44)):
v = np.asarray(direction, float)
light = pv.Light(position=tuple(-v * 100), focal_point=(0, 0, 0),
light_type='scene light')
light.intensity = intensity
self.view.add_light(light)
# ------------------------------------------------------------ actions
def _switch_part(self, part):
self.args.part = part
self.work = Workpiece(self.res, part)
self._show_edited_part()
self.statusBar().showMessage(self.work.path)
self.refresh(reset=True)
def on_add(self):
row = self.work.add_box()
self.refresh()
self.list.setCurrentRow(row)
def on_duplicate(self):
row = self.list.currentRow()
if row < 0:
return
clone = json.loads(json.dumps(self.work.elements[row]))
clone['name'] = clone.get('name', 'socket') + ' copy'
self.work.elements.append(clone)
self.work.unwrap()
self.refresh()
self.list.setCurrentRow(len(self.work.elements) - 1)
def on_remove(self):
row = self.list.currentRow()
if row >= 0:
self.work.remove(row)
self.refresh()
def on_unwrap(self):
packed = self.work.unwrap()
if packed:
self.statusBar().showMessage(
f'unwrapped onto {packed[0][0]}x{packed[0][1]}')
self.refresh()
def on_save(self):
written = self.work.save()
self.statusBar().showMessage(
'wrote ' + ', '.join(os.path.relpath(p, REPO) for p in written))
self.refresh()
def on_name(self):
row = self.list.currentRow()
if row >= 0:
self.work.elements[row]['name'] = self.name_edit.text()
self.refresh()
def on_spin(self):
row = self.list.currentRow()
if row < 0 or getattr(self, '_loading', False):
return
el = self.work.elements[row]
for key in ('from', 'to'):
el[key] = [self.spins[(key, i)].value() for i in range(3)]
self.work.unwrap()
self.refresh()
# ------------------------------------------------------------ drawing
def refresh(self, reset=False):
row = self.list.currentRow()
self._sync_list(row)
self._sync_fields(row)
self._rebuild(reset)
def _sync_list(self, row):
self._loading = True
self.list.blockSignals(True)
self.list.clear()
for i, el in enumerate(self.work.elements):
lo, hi = el['from'], el['to']
size = [round(abs(b - a), 3) for a, b in zip(lo, hi)]
self.list.addItem(f"{i}: {el.get('name', 'box')} {size}")
if 0 <= row < self.list.count():
self.list.setCurrentRow(row)
elif self.list.count():
self.list.setCurrentRow(0)
self.list.blockSignals(False)
self._loading = False
def _sync_fields(self, row):
self._loading = True
row = self.list.currentRow()
enabled = 0 <= row < len(self.work.elements)
for spin in self.spins.values():
spin.setEnabled(enabled)
self.name_edit.setEnabled(enabled)
if enabled:
el = self.work.elements[row]
self.name_edit.setText(el.get('name', 'box'))
for key in ('from', 'to'):
for i in range(3):
self.spins[(key, i)].setValue(float(el[key][i]))
size = [round(abs(b - a), 3) for a, b in zip(el['from'], el['to'])]
self.size_label.setText(' x '.join(str(s) for s in size))
else:
self.name_edit.setText('')
self.size_label.setText('-')
ts = self.work.doc.get('texture_size', [16, 16])
self.atlas_label.setText(f'{ts[0]} x {ts[1]}')
self._loading = False
def _context(self, part=None):
"""Armory's plate for a part, so new geometry has something to meet."""
spec = PARTS[part or self.args.part]
if spec['piece'] is None:
return self._icon_context()
piece = json.loads(self.res.read(PIECE_FILES[spec['piece']]))['slots']
slot = piece[spec['slot']]['transform']
plate = mc.model_quads(
self.res.model(f"miapi:models/item/armor/model/{spec['plate']}/"
'[material.texture].json', self.args.variant), self.res)
return slot, plate
def _icon_context(self):
"""The inventory sprite, for aiming icon geometry at.
The icon is a flat `item/generated` sprite, so the context here is a
picture rather than a shape - but it is the picture the gem has to land
on, and eyeballing pixel offsets against it beats counting them.
"""
icon = self.args.icon or 'miapi:models/item/armor/gui/heavy/arm_left/base/' \
'[material.texture].json'
return {}, mc.model_quads(self.res.model(icon, self.args.variant), self.res)
def _frame(self, part):
"""The matrix and pivot that put a part's model on the wearer."""
spec = PARTS[part]
if spec['piece'] is None:
return np.eye(4), (0.0, 0.0, 0.0)
piece = json.loads(self.res.read(PIECE_FILES[spec['piece']]))['slots']
return (mc.transform_matrix(piece[spec['slot']]['transform']),
mc.PIVOTS.get(part, (0.0, 0.0, 0.0)))
def _rebuild(self, reset=False):
for name in list(self.actors):
self.view.remove_actor(self.actors.pop(name))
textured = self.textured.isChecked()
groups = []
for part in PARTS:
if not self._visible(part):
continue
try:
matrix, pivot = self._frame(part)
_slot, plate = self._context(part)
except (KeyError, TypeError):
continue
groups.append((f'plate_{part}',
[q.transformed(matrix, pivot) for q in plate],
(0.60, 0.60, 0.66)))
try:
matrix, pivot = self._frame(self.args.part)
except (KeyError, TypeError):
matrix, pivot = np.eye(4), (0.0, 0.0, 0.0)
try:
work = self.work.quads(self.res)
except Exception:
work = []
groups.append(('work', [q.transformed(matrix, pivot) for q in work],
(0.95, 0.72, 0.30)))
for tag, quads, colour in groups:
if not quads:
continue
# The workpiece is coloured against its own box count, which is what
# its unwrap template was drawn against; the context plate is not
# ours and gets a flat colour so the two never look related.
count = len(self.work.elements) if tag == 'work' else None
meshes = ae.build_meshes(self.res, quads, count,
by_colour=(tag == 'work' and not textured))
for j, (mesh, tex, face_colour) in enumerate(meshes):
kw = dict(smooth_shading=False, ambient=0.42, diffuse=0.78,
specular=0.0)
if textured and tex is not None:
kw['texture'] = ae.make_texture(self.res, tex)
kw['color'] = 'white'
else:
kw['color'] = face_colour or colour
name = f'{tag}{j}'
self.actors[name] = self.view.add_mesh(mesh, name=name, **kw)
if self._visible(WEARER):
import pyvista as pv
body = mc.humanoid_body()
pts = np.concatenate([q.pts for q in body])
faces = np.hstack([[4, *range(4 * k, 4 * k + 4)] for k in range(len(body))])
self.actors['body'] = self.view.add_mesh(
pv.PolyData(pts, faces), name='body', color=(0.30, 0.32, 0.38),
opacity=0.25, smooth_shading=False, specular=0.0)
if reset:
self._aim(pivot)
if self._realised:
self.view.render()
def _aim(self, pivot):
target = (pivot[0], pivot[1] + 2.0, 0.0)
self.view.camera.position = (target[0] - 16, target[1] - 6, target[2] - 26)
self.view.camera.focal_point = target
self.view.camera.up = (0.0, -1.0, 0.0)
self.view.camera.view_angle = 34.0
self.view.camera_set = True
def main(argv=None):
ap = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument('--jar', action='append', default=[], metavar='PATH')
ap.add_argument('--part', default='left_arm', choices=sorted(PARTS))
ap.add_argument('--variant', default='default')
ap.add_argument('--icon', help='model path to show behind `item` geometry')
ap.add_argument('--background', default='#1a1a1e')
ap.add_argument('--no-textures', action='store_true')
ap.add_argument('--body', action='store_true')
args = ap.parse_args(argv)
sources = list(args.jar)
sources.append(os.path.join(REPO, 'src/main/resources'))
res = mc.Resources(sources)
app = QtWidgets.QApplication(sys.argv[:1])
window = ArmourGui(res, args)
window.resize(1420, 820)
window.show()
return app.exec()
if __name__ == '__main__':
sys.exit(main())