def process_img(image: numpy.ndarray):
# Подготавливается конфигурация
config = PredictionConfig()
config.display()
# Производится определение модели
model = MaskRCNN(mode='inference', model_dir='./', config=config)
# Производится загрузка весов
model.load_weights(
WEIGHT_FILE,
by_name=True
)
# Производим сопоставление имен идентификаторов(Поскольку обучаем на части объектов)
objs_ids = dict()
objs_names = dict()
for i, class_id in enumerate(CLASS_IDS, start=1):
# print(i)
objs_ids[i] = int(OBJECTS[class_id].split('_')[0])
objs_names[i] = OBJECTS[class_id]
print(objs_ids)
print(objs_names)
identified_objects = list()
r = model.detect([image], verbose=0)[0]
# Если на кадре найдены объекты, то проводим их обработку
if len(r['class_ids']):
for num, object_id in enumerate(r['class_ids']):
box = r['rois'][num]
identified_object = IdentifiedObject(
class_id=objs_ids[object_id],
name=objs_names[object_id],
mask=r['masks'][:, :, num],
box=Box(box[1], box[0], box[3], box[2]),
score=r['scores'][num]
)
print(identified_object.__dict__)
identified_objects.append(identified_object)
return identified_objects
def instance_segment(image_path, model_path):
# define the model
rcnn = MaskRCNN(mode='inference',
model_dir='mask_rcnn',
config=TestConfig())
# load coco model weights
# Epochs 30
# Size (3000 train and 30000 validation)\
# Steps Per Epoch 1000
rcnn.load_weights(model_path, by_name=True)
# load photograph
img = load_img(image_path)
img = img_to_array(img)
# make prediction
results = rcnn.detect([img], verbose=0)
# get dictionary for first prediction
r = results[0]
# show photo with bounding boxes, masks, class labels and scores
display_instances(img, r['rois'], r['masks'], r['class_ids'], class_names,
r['scores'])
return r