diff --git a/backend/webserver/util/mask_rcnn.py b/backend/webserver/util/mask_rcnn.py index 30efd53c..593ce072 100644 --- a/backend/webserver/util/mask_rcnn.py +++ b/backend/webserver/util/mask_rcnn.py @@ -5,6 +5,8 @@ from keras.preprocessing.image import img_to_array from mrcnn.config import Config import mrcnn.model as modellib +import logging +logger = logging.getLogger('gunicorn.error') MODEL_DIR = "/workspace/models" @@ -18,14 +20,14 @@ class CocoConfig(Config): NAME = "coco" GPU_COUNT = 1 IMAGES_PER_GPU = 1 - NUM_CLASSES = 1 + 80 + NUM_CLASSES = len(CLASS_NAMES) class MaskRCNN(): def __init__(self): - - self.config = CocoConfig() + + self.config = CocoConfig() self.model = modellib.MaskRCNN( mode="inference", model_dir=MODEL_DIR, @@ -34,26 +36,27 @@ def __init__(self): try: self.model.load_weights(COCO_MODEL_PATH, by_name=True) self.model.keras_model._make_predict_function() + logger.info(f"Loaded MaskRCNN model: {COCO_MODEL_PATH}") except: logger.error(f"Could not load MaskRCNN model (place 'mask_rcnn_coco.h5' in the models directory)") self.model = None - + def detect(self, image): if self.model is None: return {} - + image = image.convert('RGB') width, height = image.size image.thumbnail((1024, 1024)) image = img_to_array(image) result = self.model.detect([image])[0] - + masks = result.get('masks') class_ids = result.get('class_ids') - + coco_image = im.Image(width=width, height=height) for i in range(masks.shape[-1]): @@ -68,4 +71,3 @@ def detect(self, image): model = MaskRCNN() -