Had a few errors with encoding in files, and trying to read directories that were named .var
My fixes are below - sorry about formatting. Try to format it online?
#!/bin/python3
#usage:
# python3 extract_appearance.py "Saves/scene/*.json" "AddonPackages/*.var"
import json
import sys
import os
import zipfile
import re
import glob
import traceback
import logging
SCRIPT_VERSION = '2.1'
# logging.getLogger().setLevel(logging.WARNING)
logging.getLogger().setLevel(logging.INFO)
try:
from tqdm import tqdm as tqdmfunc
except ImportError:
logging.warning("tqdm not installed, not showing progress bar")
def tqdmfunc(iterator):
yield from iterator
def update_dict(d, fn, key=None):
if isinstance(d, dict):
for k, v in d.items():
d[k] = update_dict(v, fn, k)
# return not needed to update dicts by reference, but used
# here to assist in assigning subtrees
return d
elif isinstance(d, list):
return [update_dict(i, fn) for i in d]
else:
return fn(d, key)
def update_dict_selfname(d, self_name):
def update_fn(v, key=None):
if key == 'id' and isinstance(v, str):
return v.replace("SELF:/", self_name + ":/")
else:
return v
update_dict(d, update_fn)
def getPersonAtoms(scenejson):
for atom in scenejson['atoms']:
if atom['type'] == "Person":
yield atom
class BaseExtractor(object):
"""docstring for BaseExtractor"""
out_dir_verify = "./customLooks"
#out_dir_verify = "./Custom/Atom/Person/"
def __init__(self, noclobber=False):
super(BaseExtractor, self).__init__()
self.noclobber = noclobber
self.dupes = set()
self.out_dir = self.getOutDir()
os.makedirs(self.out_dir, exist_ok=True)
def getOutDir(self):
if not os.path.isdir(self.out_dir_verify):
logging.error("Script is not being run from the VaM root directory!")
logging.warning("Can't cleanly merge into Person/Appearance presets; making ExtractedAppearance folder instead.")
return self.out_dir_fallback
else:
return self.out_dir_wanted
def extractFromVar(self, var):
# Check if the path is a file and has a .var extension
if not os.path.isfile(var) or not var.endswith('.var'):
print(f"{var} is not a valid .var file or is a directory. Skipping.")
return
__, filename = os.path.split(var)
self_name, __ = os.path.splitext(filename)
author = self_name.split('.')[0]
try:
with zipfile.ZipFile(var) as varzip:
infodict = {
zi.filename: zi
for zi in varzip.infolist()
}
zip_scenes = {
k: v for k, v in infodict.items()
if k.startswith("Saves/scene") and k.endswith(".json")
}
for scenejsonpath in zip_scenes.keys():
logging.debug("%s:%s", var, scenejsonpath)
__, sjp_filename = os.path.split(scenejsonpath)
sjp_plain, sjp_ext = os.path.splitext(sjp_filename)
with varzip.open(scenejsonpath, 'r') as fp:
scenejson = json.load(fp)
if False:
update_dict_selfname(scenejson, self_name)
else:
# This method is potentially inaccurate but MUCH faster:
scenejson = json.loads(json.dumps(scenejson).replace("SELF:/", self_name + ":/"))
def readthumb():
return varzip.open(scenejsonpath.replace('.json', '.jpg'), 'r')
def outnameFn(atom):
return f"Preset_{self_name}.{sjp_plain}{atom['id'].replace('Person', '')}".replace('#', '').replace('/', '-')
self.extractFromSceneJson(scenejson, outnameFn, readthumb)
except json.JSONDecodeError:
print(var)
traceback.print_exc()
os.rename(var, var + '.invalid')
except zipfile.BadZipFile:
print(var)
traceback.print_exc()
os.rename(var, var + '.invalid')
except Exception as e:
print(f"An error occurred with file {var}: {e}. Skipping to next file.")
traceback.print_exc()
def extractFromSceneJsonPath(self, scenejsonpath):
__, sjp_filename = os.path.split(scenejsonpath)
sjp_plain, sjp_ext = os.path.splitext(sjp_filename)
print('loading: ' + scenejsonpath)
def readthumb():
return open(scenejsonpath.replace('.json', '.jpg'), 'rb')
def outnameFn(atom):
return f"Preset_!local.{sjp_plain}{atom['id'].replace('Person', '')}".replace('#', '').replace('/', '-')
try:
with open(scenejsonpath, 'r', encoding='utf-8') as fp:
self.extractFromSceneJson(json.load(fp), outnameFn, readthumb)
except json.JSONDecodeError:
print(f"Error decoding JSON in file {scenejsonpath}. Skipping to next file.")
except Exception as e:
print(f"An error occurred with file {scenejsonpath}: {e}. Skipping to next file.")
def presetOutname(self, outname, ext):
return os.path.join(self.out_dir, outname + ext)
def extractFromSceneJson(self, scenejson, outnameFn, readthumb):
# print([atom['id'] for atom in scenejson['atoms'] if atom['type'] == 'Person'])
for atom in getPersonAtoms(scenejson):
outname = outnameFn(atom)
if self.noclobber and os.path.isfile(self.presetOutname(outname, '.vap')):
logging.info(f"Skipping {outname!r}")
continue
# print(atom['id'])
preset = {
"setUnlistedParamsToDefault": "true",
'storables': self.filterPersonStorables(atom)
}
self.savePreset(preset, outname, readthumb)
def savePreset(self, preset, outname, readthumb):
preset_hash = json.dumps(preset, sort_keys=True)
if preset_hash not in self.dupes:
outpath = self.presetOutname(outname, '.vap')
with open(outpath, 'w', encoding='utf-8') as fp:
logging.info("-> %s", outpath)
fp.write(json.dumps(preset, indent=3, ensure_ascii=False))
try:
img_outpath = self.presetOutname(outname, '.jpg')
with open(img_outpath, 'wb') as fp:
with readthumb() as fp2:
fp.write(fp2.read())
except (KeyError, FileNotFoundError):
# No image in archive
pass
self.dupes.add(preset_hash)
else:
# logging.warn("Skip duplicate %s", outname)
pass
def filterPersonStorables(self, atom):
raise NotImplementedError()
class AppearanceExtractor(BaseExtractor):
out_dir_wanted = "./Custom/Atom/Person/Appearance/extracted/"
out_dir_fallback = "./ExtractedAppearance/"
def filterPersonStorables(self, atom):
storables = atom['storables']
storables_new = []
for storable in storables:
storable.pop("position", None)
storable.pop("rotation", None)
# Remove all keys containing these substrings
for key in [*storable.keys()]:
if any(ss in key.lower() for ss in ['position', 'rotation']):
storable.pop(key)
# Remove all storables whose ids contain these substrings
if any(ss in storable['id'].lower() for ss in ["control", "trigger", "plugin", "preset", "animation"]):
# print("pop", storable['id'])
storables.remove(storable)
continue
# Remove transient morphs
if storable['id'] == "geometry" and 'morphs' in storable:
for morph in [*storable['morphs']]:
if any(re.match(ss, morph.get('uid', '')) for ss in [
r'Breast Impact',
r'^OpenXXL$',
r'^Eyelids (Top|Bottom) (Down|Up) (Left|Right)$', r'Brow .*(Up|Down)', r'^(Left|Right) Fingers',
r'^Mouth Open', r'Tongue In-Out', r'^Smile',
'Shock', 'Surprise', 'Fear', 'Pain', 'Concentrate', 'Eyes Closed', r'^Flirting',
]):
# print(morph)
storable['morphs'].remove(morph)
# Remove storables with just an id and no properties
if len(storable.keys()) > 1:
storables_new.append(storable)
return storables_new
class OutfitExtractor(BaseExtractor):
out_dir_wanted = "./Custom/Atom/Person/Clothing/extracted/"
out_dir_fallback = "./ExtractedOutfits/"
def filterPersonStorables(self, atom):
storables = atom['storables']
clothing_master = next(filter(lambda v: v['id'] == 'geometry' and 'clothing' in v.keys(), storables))
clothing = clothing_master['clothing']
try:
clothing_ids = [c.get('internalId', c['id']) for c in clothing]
except KeyError:
print(clothing)
raise
storables_new = [
{
"id": "geometry",
"clothing": clothing
}
]
for storable in storables:
if any(storable['id'].startswith(id_) for id_ in clothing_ids):
storables_new.append(storable)
# else:
# print(storable['id'], clothing_ids)
# print(storables_new)
return storables_new
def add_bool_arg(parser, name, default=False, **kwargs):
group = parser.add_mutually_exclusive_group(required=False)
group.add_argument('--' + name, dest=name, action='store_true', **kwargs)
group.add_argument('--no-' + name, dest=name, action='store_false')
parser.set_defaults(**{name: default})
def parse_args():
import argparse
parser = argparse.ArgumentParser(
description="Exports presets",
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument('input_globs', help="Input files or fileglobs", nargs='+')
add_bool_arg(parser, 'appearance', default=True, help="Extract appearance (body) presets")
add_bool_arg(parser, 'outfit', default=False, help="Extract outfit (clothing) presets")
parser.add_argument('--no-clobber', action='store_true', help="If set, don't overwrite files.", default=False)
return parser.parse_args()
def main():
print('Appearance Preset Extractor', SCRIPT_VERSION, 'by cartitolaire')
if not sys.argv[1:]:
print('')
print('This tool automatically exports presets from existing scenes and collects them in an "extracted" folder under presets. The presets will use the thumbnail for the source scene.')
print('')
print("Example usages:")
print(' ', sys.argv[0], 'AddonPackages/Anon.SceneName.1.var')
print(' ', sys.argv[0], '"AddonPackages/Anon*.var"')
print(' ', sys.argv[0], '"Saves/scene/*.json" "AddonPackages/*.var"')
print(' ', sys.argv[0], '--outfit --no-appearance AddonPackages/Anon.SceneName.1.var')
return
args = parse_args()
extractors = []
if args.appearance:
extractors.append(AppearanceExtractor)
if args.outfit:
extractors.append(OutfitExtractor)
try:
for etype in extractors:
extractor = etype(noclobber=args.no_clobber)
iterable = sum((glob.glob(a, recursive=True) for a in args.input_globs), [])
for filepath in tqdmfunc(iterable):
__, filename = os.path.split(filepath)
self_name, ext = os.path.splitext(filename)
try:
if ext == ".var":
extractor.extractFromVar(filepath)
elif ext == ".json":
extractor.extractFromSceneJsonPath(filepath)
else:
raise NotImplementedError(ext)
except KeyError:
pass
except KeyboardInterrupt:
return
if __name__ == "__main__":
main()
