Skip to content

Commit d757cc5

Browse files
luxedojsbroks
authored andcommitted
Fixed Mask-RCNN model logging and NUM_CLASSES (#201)
1 parent 538ed1a commit d757cc5

File tree

1 file changed

+10
-8
lines changed

1 file changed

+10
-8
lines changed

backend/webserver/util/mask_rcnn.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
from keras.preprocessing.image import img_to_array
66
from mrcnn.config import Config
77
import mrcnn.model as modellib
8+
import logging
9+
logger = logging.getLogger('gunicorn.error')
810

911

1012
MODEL_DIR = "/workspace/models"
@@ -18,14 +20,14 @@ class CocoConfig(Config):
1820
NAME = "coco"
1921
GPU_COUNT = 1
2022
IMAGES_PER_GPU = 1
21-
NUM_CLASSES = 1 + 80
23+
NUM_CLASSES = len(CLASS_NAMES)
2224

2325

2426
class MaskRCNN():
2527

2628
def __init__(self):
27-
28-
self.config = CocoConfig()
29+
30+
self.config = CocoConfig()
2931
self.model = modellib.MaskRCNN(
3032
mode="inference",
3133
model_dir=MODEL_DIR,
@@ -34,26 +36,27 @@ def __init__(self):
3436
try:
3537
self.model.load_weights(COCO_MODEL_PATH, by_name=True)
3638
self.model.keras_model._make_predict_function()
39+
logger.info(f"Loaded MaskRCNN model: {COCO_MODEL_PATH}")
3740
except:
3841
logger.error(f"Could not load MaskRCNN model (place 'mask_rcnn_coco.h5' in the models directory)")
3942
self.model = None
40-
43+
4144

4245
def detect(self, image):
4346

4447
if self.model is None:
4548
return {}
46-
49+
4750
image = image.convert('RGB')
4851
width, height = image.size
4952
image.thumbnail((1024, 1024))
5053

5154
image = img_to_array(image)
5255
result = self.model.detect([image])[0]
53-
56+
5457
masks = result.get('masks')
5558
class_ids = result.get('class_ids')
56-
59+
5760
coco_image = im.Image(width=width, height=height)
5861

5962
for i in range(masks.shape[-1]):
@@ -68,4 +71,3 @@ def detect(self, image):
6871

6972

7073
model = MaskRCNN()
71-

0 commit comments

Comments
 (0)