55from keras .preprocessing .image import img_to_array
66from mrcnn .config import Config
77import mrcnn .model as modellib
8+ import logging
9+ logger = logging .getLogger ('gunicorn.error' )
810
911
1012MODEL_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
2426class 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
7073model = MaskRCNN ()
71-
0 commit comments