pythonflaskfaster-rcnn

Running a mask-rcnn model on Flask server startup


This is my current Flask code that works fine, it receives a POST request with the image from the client, runs it through the model (based on this GH: https://github.com/matterport/Mask_RCNN), and sends a masked image back to the client.

However, it is loading the model from the Configuration file and loading the weights for each request, which takes ages. I want to load the model on server startup and the weights and pass that to the index function. I have tried the solutions from other questions, but with no luck. I wonder if it's because I am loading a model, and then weights, rather than just loading a single h5 model file?

How do I load a file on initialization in a flask application Run code after flask application has started

Flask app:

from flask import Flask, jsonify, request
import base64
import cv2
import numpy as np
from Configuration import create_model


app = Flask(__name__)

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == "POST":
        # Load the image sent from the client
        imagefile = request.files['image'].read()  # Type: bytes
        jpg_as_np = np.frombuffer(imagefile, dtype=np.uint8) # Convert to numpy array
        img = cv2.imdecode(jpg_as_np, flags=1) # Decode from numpy array to opencv object - This is an array

        ### Enter OpenCV/Tensorflow below ###

        model = create_model()

        image = img[..., ::-1]

        # Detect objects
        r = model.detect([image], verbose=0)[0]

        REDACTED VISUALIATION CODE

        ### ###

        string = base64.b64encode(cv2.imencode('.jpg', masked_image)[1]).decode() # Convert back to b64 string ready for json.

        return jsonify({"count": str(r["masks"].shape[2]), 'image': string})

if __name__ == "__main__":
    app.run()

Configuration:

def create_model():
    device = "/cpu:0"
    weights_path = "weights.h5"
    with tf.device(device):
        model = modellib.MaskRCNN(mode="inference", model_dir=weights_path, config=InferenceConfig())
    model.load_weights(weights_path, by_name=True)
    print("Weights Loaded")
    return model

Solution

  • I solved this using the before_first_request decorator. Below is the general structure:

    app = Flask(__name__)
    
    @app.before_first_request
    def before_first_request_func():
        MOODEL WEIGHT LOADING CODE
        return model
    
    @app.route('/', methods=['POST'])
    def index():
        if request.method == "POST":
            REDACTED LOADING CODE
            # Detect objects
            r = model.detect([image], verbose=0)[0]
    
            REDACTED VISUALISATION CODE
    
            string = base64.b64encode(cv2.imencode('.jpg', masked_image)[1]).decode() # Convert back to b64 string ready for json.
    
            return jsonify({"count": str(r["masks"].shape[2]), 'image': string})
    
    if __name__ == "__main__":
        app.run()
    

    model is stored in memory and can be referenced inside the detection function later on. It is usable for each POST request and does not need to be reloaded.