Files
dragonpilot/selfdrive/modeld/thneed/lib.py
dragonpilot 1542fedcd1 dragonpilot beta3
date: 2023-08-11T23:00:23
commit: af14371afdf5023a178ed6f21a82099ea64dffe4
2023-08-11 23:00:38 +08:00

33 lines
859 B
Python

import struct
import json
def load_thneed(fn):
with open(fn, "rb") as f:
json_len = struct.unpack("I", f.read(4))[0]
jdat = json.loads(f.read(json_len).decode('latin_1'))
weights = f.read()
ptr = 0
for o in jdat['objects']:
if o['needs_load']:
nptr = ptr + o['size']
o['data'] = weights[ptr:nptr]
ptr = nptr
for o in jdat['binaries']:
nptr = ptr + o['length']
o['data'] = weights[ptr:nptr]
ptr = nptr
return jdat
def save_thneed(jdat, fn):
new_weights = []
for o in jdat['objects'] + jdat['binaries']:
if 'data' in o:
new_weights.append(o['data'])
del o['data']
new_weights_bytes = b''.join(new_weights)
with open(fn, "wb") as f:
j = json.dumps(jdat, ensure_ascii=False).encode('latin_1')
f.write(struct.pack("I", len(j)))
f.write(j)
f.write(new_weights_bytes)