schemagen/schemgen.py

132 lines
4.6 KiB
Python
Raw Normal View History

2021-09-30 13:18:38 +08:00
from PIL import Image
from pathlib import Path
from itertools import zip_longest
import re
def read_palette(text):
palette = dict()
metadata = dict()
in_body = False
for line in text.split('\n'):
line = line.strip()
if not line or line == "GIMP Palette":
continue
if line == '#':
in_body = True
continue
if in_body:
*rgb, name = line.split()
r, g, b = [int(x) for x in rgb]
index = len(palette)
palette[f"{r:02x}{g:02x}{b:02x}"] = {"rgb": (r, g, b), "name": name if name.startswith("minecraft:") else f"minecraft:{name}", "index": index}
else:
k, *v = line.split(':')
k, v = k.strip(), ':'.join(v).strip()
metadata[k] = v
return {"metadata": metadata, "palette": palette}
def read_imtuple(impath, palette_path):
palette = read_palette(Path(palette_path).read_text())
im = Image.open(impath)
im = im.rotate(180)
im = im.convert('RGB')
used_color_hex = tuple(set(f"{r:02x}{g:02x}{b:02x}" for r, g, b in im.getdata()))
used_mc_palette = tuple(palette["palette"][h]["name"] for h in used_color_hex)
imdata = (used_color_hex.index(f"{r:02x}{g:02x}{b:02x}") for r, g, b in im.getdata())
#imtuple = tuple(zip_longest(*([iter(imdata)] * im.width), fillvalue=0))
#assert len(imtuple) == im.height
imtuple = imdata
return (used_mc_palette, imtuple, im.width, im.height)
def read_text(text):
from mcdraw import text_to_im
im = text_to_im(text)
used_mc_palette = ('minecraft:air', 'minecraft:stone')
imdata = (0 if i else 1 for i in im.getdata())
imtuple = imdata
return (used_mc_palette, imtuple, im.width, im.height)
import gzip
from pynbt import NBTFile, TAG_Int, TAG_Compound, TAG_Short, TAG_Byte_Array, TAG_List, TAG_Int_Array
from io import BytesIO
import struct
class VarInt:
MAX_BYTES = 5
def __init__(self, io, set_value=None):
self._io = io
if set_value:
self.setval(set_value)
def setval(self, value):
ret = bytes()
while True:
byte = value & 0x7F
value >>= 7
ret += struct.pack("B", byte | (0x80 if value > 0 else 0))
if value == 0:
break
if len(ret) > self.MAX_BYTES:
raise ValueError("Tried to write too long of a VarInt")
self._io.write(ret)
@property
def value(self):
number = 0
bytes_encountered = 0
while True:
byte = self._io.read(1)
if len(byte) < 1:
raise EOFError("Unexpected end of message.")
byteord = ord(byte)
number |= (byteord & 0x7F) << 7 * bytes_encountered
if not byteord & 0x80:
break
bytes_encountered += 1
if bytes_encountered > self.MAX_BYTES:
raise ValueError("Tried to read too long of a VarInt")
return number
def assemble_nbtfile(used_mc_palette, imtuple, width, height):
blackdataio = BytesIO()
varint = VarInt(blackdataio)
for i in imtuple:
varint.setval(i)
nbtvalue = {
'Version': TAG_Int(2),
'PaletteMax': TAG_Int(len(used_mc_palette)),
'Palette': TAG_Compound(
{name: TAG_Int(i) for i, name in enumerate(used_mc_palette)}
),
'Metadata': TAG_Compound({
'WEOffsetX': TAG_Int(0),
'WEOffsetY': TAG_Int(0),
'WEOffsetZ': TAG_Int(0)
}),
'Width': TAG_Short(width),
'Height': TAG_Short(1),
'Length': TAG_Short(height),
'DataVersion': TAG_Int(2586),
'BlockData': TAG_Byte_Array(blackdataio.getvalue()),
'BlockEntities': TAG_List(TAG_Compound),
'Offset': TAG_Int_Array([0, 0, 0])
}
nbt = NBTFile(value=nbtvalue, name='Schematic')
return nbt
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description='schemgen')
parser.add_argument('-o', '--output', default='output.schem', help='output worldedit schema file')
parser.add_argument('-i', '--input', type=str, default="input.png", help='input image')
parser.add_argument('-t', '--text', type=str, default="", help='input text')
parser.add_argument('-p', '--palette', type=str, default="minecraft.gpl", help='gimp palette file')
args = parser.parse_args()
if args.text:
nbt = assemble_nbtfile(*read_text(args.text))
else:
nbt = assemble_nbtfile(*read_imtuple(args.input, args.palette))
print(nbt.pretty())
with gzip.open(args.output, 'w') as f:
nbt.save(f)