Pytorch 约罗和普罗托夫

我正在使用RedisAI构建一个视频处理管道。我使用的模板加载TinyYOLO protobuf，然后加载PyTorch脚本：

print('Loading model - ', end='')
with open('models/tiny-yolo-voc.pb', 'rb') as f:
    model = f.read()
    res = conn.execute_command('AI.MODELSET', 'yolo:model', 'TF', args.device, 'INPUTS', 'input', 'OUTPUTS', 'output', model)
    print(res)

# Load the PyTorch post processing boxes script
print('Loading script - ', end='')
with open('yolo_boxes.py', 'rb') as f:
    script = f.read()
    res = conn.execute_command('AI.SCRIPTSET', 'yolo:script', args.device, script)
    print(res)

然后，它在每一帧上运行RedisAI模型运行程序，生成一些框，并抛出类别不是14（人）的框

我的问题是，如果我想使用YOLO4或YOLO5，我是否能够接受一个经过训练的模型，从中生成一个protobuf，然后将其放入这个过程中？类别的编号是否会有所不同

modelRunner = redisAI.createModelRunner('yolo:model')
    redisAI.modelRunnerAddInput(modelRunner, 'input', image_tensor)
    redisAI.modelRunnerAddOutput(modelRunner, 'output')
    model_replies = redisAI.modelRunnerRun(modelRunner)
    model_output = model_replies[0]
    prf.add('model')

# log('script')
# The model's output is processed with a PyTorch script for non maxima suppression
scriptRunner = redisAI.createScriptRunner('yolo:script', 'boxes_from_tf')
redisAI.scriptRunnerAddInput(scriptRunner, model_output)
redisAI.scriptRunnerAddOutput(scriptRunner)
script_reply = redisAI.scriptRunnerRun(scriptRunner)
prf.add('script')

# log('boxes')
# The script outputs bounding boxes
shape = redisAI.tensorGetDims(script_reply)
buf = redisAI.tensorGetDataAsBlob(script_reply)
boxes = np.frombuffer(buf, dtype=np.float32).reshape(shape)

# Iterate boxes to extract the people
ratio = float(IMG_SIZE) / max(pil_image.width, pil_image.height)  # ratio = old / new
pad_x = (IMG_SIZE - pil_image.width * ratio) / 2                  # Width padding
pad_y = (IMG_SIZE - pil_image.height * ratio) / 2                 # Height padding
boxes_out = []
people_count = 0
for box in boxes[0]:
    if box[4] == 0.0:  # Remove zero-confidence detections
        continue
    if box[-1] != 14:  # Ignore detections that aren't people
        continue
    people_count += 1