Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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='resources/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)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement