abse/museair-usb

scripts/package-native-assembly-pcb.py

"""Package a native CadViewer export as ONE PCB CAD item with native textures and lit board-face shading.
Usage: python scripts/package-native-assembly-pcb.py /path/native-pcb.glb /path/textures
External buffer chunks avoid the registry's per-file upload limit; they are still one model.
"""
import base64, copy, hashlib, io, json, struct, sys, urllib.parse
from pathlib import Path
import numpy as np
from PIL import Image
r = Path(__file__).resolve().parent.parent
src = Path(sys.argv[1]); textures = Path(sys.argv[2]); assets = r/'assembly-assets'
b = src.read_bytes(); n = struct.unpack_from('<I', b, 12)[0]
j = json.loads(b[20:20+n]); raw = b[28+n:]
# A non-indexed multi-material export can paint every face with every material.
# Reject those corrupt exports instead of hiding the problem by recoloring CAD.
for m in j['meshes']:
    primitives=m['primitives']
    if len(primitives)>1:
        assert all('indices' in p for p in primitives), 'Re-export indexed multi-material geometry before packaging'
        assert len({p['indices'] for p in primitives})==len(primitives), 'Duplicated whole-mesh material groups'
index = json.loads((r/'dist/index/circuit.json').read_text())
refs = {x['source_component_id']:x['name'] for x in index if x['type']=='source_component'}
names = {x['cad_component_id']:refs.get(x['source_component_id'],'') for x in index if x['type']=='cad_component'}
seen=set()
def rename(i, parent='PCB'):
    node=j['nodes'][i]; name=names.get(node.get('extras',{}).get('cad_component_id'),parent)
    if name!='PCB':seen.add(name)
    if 'mesh' in node:node['name']=f'{name}_{i}'
    for child in node.get('children',[]):rename(child,name)
for scene in j['scenes']:
    for i in scene['nodes']:rename(i)
assert len(seen)==65 and 'U3' not in seen, seen
# Remove per-part CAD IDs from embedded descendants so picking resolves to the
# single assembly wrapper, rather than an unrelated assembly component ID.
for node in j['nodes']:
    for key in ['cad_component_id','pcb_component_id']:
        node.get('extras',{}).pop(key,None)
# Native viewer export already uses PCB XY / Z-up. Preserve its original scene graph.
pose=[0,1,0,0,-1,0,0,0,0,0,1,0,2.5,0,37.8,1]
for scene in j['scenes']:
    j['nodes'].append({'name':'PCB_COMPLETE','matrix':pose,'children':scene['nodes']})
    scene['nodes']=[len(j['nodes'])-1]
def glb(doc, data):
    header=json.dumps(doc,separators=(',',':')).encode();header+=b' '*((-len(header))%4)
    data+=b'\0'*((-len(data))%4)
    return struct.pack('<III',0x46546c67,2,28+len(header)+len(data))+struct.pack('<II',len(header),0x4e4f534a)+header+struct.pack('<II',len(data),0x004e4942)+data
# Keep the original pixel/geometry data embedded in the local audit model.
world=glb(copy.deepcopy(j),raw)
(r/'mechanical/horizontal/models/pcb.glb').write_bytes(world)
(r/'mechanical/horizontal/reference/pcb-native-world.glb').write_bytes(world)
rid=json.loads((r/'validation/assembly-pcb-quality-publication.json').read_text())['package_release_id']
def url(path):
    return 'https://registry-api.tscircuit.com/package_files/download?'+urllib.parse.urlencode({'package_release_id':rid,'file_path':'assembly-assets/'+path})
# Preserve native lit materials, including board clearcoat and connector metal.
# Only make textured board faces double-sided for imported scene transforms.
for mat in j['materials']:
    if 'baseColorTexture' in mat.get('pbrMetallicRoughness',{}):
        mat['doubleSided']=True
image_views={im['bufferView'] for im in j['images']}; mapping={}; views=[]; chunks=[bytearray()]
for old,v in enumerate(j['bufferViews']):
    if old in image_views:continue
    start=v.get('byteOffset',0); data=raw[start:start+v['byteLength']]
    assert len(data)<2500000
    pad=(-len(chunks[-1]))%4
    if len(chunks[-1])+pad+len(data)>2500000:chunks.append(bytearray());pad=0
    chunks[-1].extend(b'\0'*pad)
    new=dict(v,buffer=len(chunks)-1,byteOffset=len(chunks[-1]));chunks[-1].extend(data)
    mapping[old]=len(views);views.append(new)
for a in j['accessors']:
    if 'bufferView' in a:a['bufferView']=mapping[a['bufferView']]
    if 'sparse' in a:
        for key in ['indices','values']:a['sparse'][key]['bufferView']=mapping[a['sparse'][key]['bufferView']]
j['bufferViews']=views;j['buffers']=[]
files=[]
for i,data in enumerate(chunks):
    name=f'pcb-indexed-buffer-{i}-{hashlib.sha256(data).hexdigest()[:16]}.bin';(assets/name).write_bytes(data);files.append(name)
    j['buffers'].append({'uri':url(name),'byteLength':len(data)})
texture_stats=[]
for i,im in enumerate(j['images']):
    v=json.loads(b[20:20+n])['bufferViews'][im['bufferView']]
    original=Image.open(io.BytesIO(raw[v.get('byteOffset',0):v.get('byteOffset',0)+v['byteLength']])).convert('RGBA')
    data=(textures/f'texture-{i}.webp').read_bytes();decoded=Image.open(io.BytesIO(data)).convert('RGBA')
    # WebP may discard invisible RGB, but every visible channel and all alpha must match.
    a=np.array(original);d=np.array(decoded)
    assert np.array_equal(a[:,:,3],d[:,:,3]) and np.array_equal(a[a[:,:,3]>0],d[a[:,:,3]>0])
    assert len(data)+2048<4194304,(i,len(data)) # Larger images use multipart upload; no pixel loss
    name=f'pcb-texture-{i}-{hashlib.sha256(data).hexdigest()[:16]}.webp';(assets/name).write_bytes(data);files.append(name)
    j['images'][i]={'uri':url(name),'mimeType':'image/webp'}
    texture_stats.append({'image':i,'resolution':list(original.size),'bytes':len(data),'visible_pixels_lossless':True})
for tex in j['textures']:
    tex.setdefault('extensions',{})['EXT_texture_webp']={'source':tex.pop('source')}
for key in ['extensionsUsed','extensionsRequired']:
    j.setdefault(key,[]).append('EXT_texture_webp')
for key in ['extensionsUsed','extensionsRequired']:
    j[key]=list(dict.fromkeys(j.get(key,[])))
model=json.dumps(j,separators=(',',':'));(assets/'pcb.gltf').write_text(model)
(assets/'pcb.ts').write_text('export default '+json.dumps('data:model/gltf+json;base64,'+base64.b64encode(model.encode()).decode())+'\n')
rep={'status':'pass','origin':'actual native @tscircuit/3d-viewer scene of current index','source_sha256':hashlib.sha256(src.read_bytes()).hexdigest(),'source_index_sha256':hashlib.sha256((r/'dist/index/circuit.json').read_bytes()).hexdigest(),'pcb_component_models':len(seen),'component_refs':sorted(seen),'single_viewer_item':'PCB_COMPLETE','component_materials_preserved':True,'indexed_material_groups_verified':True,'native_board_texture_colors_preserved':True,'board_face_rendering':'native PBR materials including original clearcoat/specular; double-sided textured board faces','native_geometry_preserved':True,'textures':texture_stats,'external_assets':files,'buffer_bytes':sum(map(len,chunks)),'world_transform':pose,'sensor_excluded':True}
(r/'validation/assembly-native-pcb.json').write_text(json.dumps(rep,indent=2)+'\n')
print(json.dumps({'components':len(seen),'buffer_chunks':len(chunks),'model_bytes':len(model),'textures':texture_stats}))