def infer_on_image(image, graph, model, idx_to_class_title, project_meta, confidence_tag_meta, min_confidence=0): with graph.as_default(): [results] = model.detect([image], verbose=0) res_labels = [] for mask_idx, class_id in enumerate(results['class_ids']): confidence = results['scores'][mask_idx] if confidence < min_confidence: continue bool_mask = results['masks'][:, :, mask_idx] != 0 class_geometry = Bitmap(data=bool_mask) cls_title = idx_to_class_title[class_id] label = Label(geometry=class_geometry, obj_class=project_meta.get_obj_class(cls_title)) confidence_tag = Tag(confidence_tag_meta, value=round(float(confidence), 4)) label = label.add_tag(confidence_tag) res_labels.append(label) return res_labels
def detection_preds_to_sly_rects( idx_to_class, network_prediction: DetectionNetworkPrediction, img_shape, min_score_threshold, score_tag_meta) -> list: """ Converts network detection results to Supervisely Labels with Rectangle geometry. Args: idx_to_class: Dict matching predicted boxes with appropriate ObjClass. network_prediction: Network predictions packed into DetectionNetworkPrediction instance. img_shape: Size(height, width) of image that was used for inference. min_score_threshold: All detections with less scores will be dropped. score_tag_meta: TagMeta instance for score tags. Returns: A list containing labels with detection rectangles. """ labels = [] thr_mask = np.squeeze(network_prediction.scores) > min_score_threshold for box, class_id, score in zip( np.squeeze(network_prediction.boxes)[thr_mask], np.squeeze(network_prediction.classes)[thr_mask], np.squeeze(network_prediction.scores)[thr_mask]): xmin = round(float(box[1] * img_shape[1])) ymin = round(float(box[0] * img_shape[0])) xmax = round(float(box[3] * img_shape[1])) ymax = round(float(box[2] * img_shape[0])) rect = Rectangle(top=ymin, left=xmin, bottom=ymax, right=xmax) class_obj = idx_to_class[int(class_id)] label = Label(geometry=rect, obj_class=class_obj) score_tag = Tag(score_tag_meta, value=round(float(score), 4)) label = label.add_tag(score_tag) labels.append(label) return labels