def get_prediction(image, info, display=True): ''' predict on the model and return the prediction containing masks and bounding boxes. ''' vis = Visualizer() info["transforms"] = build_transforms(info["cfg"], min_image_size=min(image.shape[:-1]), max_image_size=max(image.shape[:-1])) prediction = vis.compute_prediction(info["cfg"], image, info["transforms"], info["device"], info["model"], info["cpu_device"], info["masker"]) prediction = vis.select_top_predictions(prediction, info["confidence_threshold"]) print(prediction) if display: vis.display_instances(image[:, :, ::-1], prediction.bbox, prediction.get_field("mask"), prediction.get_field("labels"), info["CATEGORIES"], prediction.get_field("scores"), show_mask=False) return prediction