import torch
import skyhook.pytorch.model as model
import skyhook.pytorch.util as util
in_dataset_id = int(sys.argv[1])
device = torch.device('cuda:0')
#device = torch.device('cpu')
model_path = 'data/items/{}/model.pt'.format(in_dataset_id)
save_dict = torch.load(model_path)
example_inputs = save_dict['example_inputs']
util.inputs_to_device(example_inputs, device)
net = model.Net(save_dict['arch'],
save_dict['comps'],
example_inputs,
save_dict['example_metadatas'],
infer=True,
device=device)
net.to(device)
net.load_state_dict(save_dict['model'])
net.eval()
stdin = sys.stdin.detach()
while True:
header = stdin.read(8)
if not header:
break
left_count, right_count = struct.unpack('>II', header)
buf = stdin.read(left_count * 64 * 64 * 3)
left_arr = numpy.frombuffer(buf, dtype='uint8').reshape(
comp_params = json.loads(arch['Components'][0].get('Params', '{}'))
comp_params['mode'] = mode
arch['Components'][0]['Params'] = json.dumps(comp_params)
# example inputs
im_data = numpy.zeros((416, 416, 3), dtype='uint8')
example_inputs = [
util.collate('image', [util.prepare_input('image', im_data, {}, {})]),
util.collate(
'detection',
[util.prepare_input('detection', [], {'CanvasDims': [416, 416]}, {})]),
]
util.inputs_to_device(example_inputs, device)
# example metadata
with open(os.path.join(yolo_path, 'data', 'coco.yaml'), 'r') as f:
d = yaml.load(f, Loader=yaml.FullLoader)
categories = d['names']
example_metadatas = [{}, {'Categories': categories}]
net = model.Net(arch, comps, example_inputs, example_metadatas, device=device)
sys.path.append(yolo_path)
orig_dict = torch.load(in_fname)['model'].state_dict()
state_dict = {}
for k, v in orig_dict.items():
state_dict['mlist.0.model.' + k] = v
net.load_state_dict(state_dict)
torch.save(net.get_save_dict(), out_fname)
}
for comp_idx, comp_spec in enumerate(arch['Components']):
comp_params = {}
if comp_spec['Params']:
comp_params = json.loads(comp_spec['Params'])
if overwrite_comp_params.get(comp_idx, None):
comp_params.update(json.loads(overwrite_comp_params[comp_idx]))
comp_spec['Params'] = json.dumps(comp_params)
example_inputs = save_dict['example_inputs']
util.inputs_to_device(example_inputs, device)
net = model.Net(arch,
save_dict['comps'],
example_inputs,
save_dict['example_metadatas'],
output_datasets=params['OutputDatasets'],
infer=True,
device=device)
net.to(device)
net.load_state_dict(save_dict['model'])
net.eval()
input_options = {}
for spec in params['InputOptions']:
input_options[spec['Idx']] = json.loads(spec['Value'])
meta = None