from mrcnn.model import MaskRCNN
from mrcnn.config import Config
import numpy as np
import cv2

class_list = [
        'large_bolt',
        'short_bolt',
        'main_housing',
        'bearing_sealing',
        'top_lid',
        'cup',
        'cylinder_flat',
        'cylinder_long',
        'cylinder_medium',
        'cylinder_short',
        'rotor_female',
        'rotor_male',
        'washer',
    ]

class PredictionConfig(Config):
    NAME = "prediction_cfg"
    NUM_CLASSES = len(class_list) + 1
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    RPN_ANCHOR_RATIOS = [0.55, 1.0, 1.57]

model_file = 'MaskRCNN_C40.h5'

config = PredictionConfig()
model = MaskRCNN(mode='inference', config=config, model_dir='/logs')
model.load_weights(model_file, by_name=True)

image = cv2.imread('test.jpg')
sample = np.expand_dims(image, 0)
result = model.detect(sample, verbose=0)[0]

detections = []
for roi, class_id, score in zip(result['rois'], result['class_ids'], result['scores']):
    class_name = class_list[class_id - 1]
    print("Class:", class_name)
    print("ROI:")
    print(roi)
    print("Score:", score)