Caffe2 model example

main.py

from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace
from detectron2.export.caffe2_inference import ProtobufDetectionModel
from detectron2.export.api import Caffe2Model
import numpy as np
import os
import torch
import cv2
print("Required modules imported.")

def create_caffe2_model():
    predict_net = caffe2_pb2.NetDef()
    with open("caffe2_model/model.pb", 'rb') as f:
        predict_net.ParseFromString(f.read())

    init_net = caffe2_pb2.NetDef()
    with open("caffe2_model/model_init.pb", 'rb') as f:
        init_net.ParseFromString(f.read())

    model = ProtobufDetectionModel(predict_net=predict_net, init_net=init_net)

    return model

def detect_people(path, model):

    img = cv2.imread(path)
    img = cv2.resize(img, (144, 144))
    img = img.swapaxes(1, 2).swapaxes(0, 1)
    results = model([{'image': torch.Tensor(img)}])

    for index_box, box in enumerate(results[0]['instances'].pred_boxes.tensor.numpy()):
        confidence = results[0]['instances'].scores.numpy()[index_box]
        classes = results[0]['instances'].pred_classes.numpy()
        if 0 in classes and confidence > 0.6:
            return True
        return False

if __name__ == "__main__": 
    model = create_caffe2_model()
    print(detect_people('images/frame_fcb3736d-c1c3-4326-817c-66b00539e19f.jpg', model))