def predict(self, src_image): param = self.getParam() # Initialize init_logging() half = self.device.type != 'cpu' # half precision only supported on CUDA # Load model if self.model is None or param.update: self.model = attempt_load(param.model_path, map_location=self.device) # load FP32 model stride = int(self.model.stride.max()) # model stride param.input_size = check_img_size(param.input_size, s=stride) # check img_size if half: self.model.half() # to FP16F # Get names and colors self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names] param.update = False else: stride = int(self.model.stride.max()) # model stride # Resize image image = letterbox(src_image, param.input_size, stride)[0] image = image.transpose(2, 0, 1) image = np.ascontiguousarray(image) self.emitStepProgress() # Run inference image = torch.from_numpy(image).to(self.device) image = image.half() if half else image.float() # uint8 to fp16/32 image /= 255.0 # 0 - 255 to 0.0 - 1.0 if image.ndimension() == 3: image = image.unsqueeze(0) self.emitStepProgress() # Inference pred = self.model(image, augment=param.augment)[0] self.emitStepProgress() # Apply NMS pred = non_max_suppression(pred, param.conf_thres, param.iou_thres, agnostic=param.agnostic_nms) self.emitStepProgress() graphics_output = self.getOutput(1) graphics_output.setNewLayer("YoloV5") graphics_output.setImageIndex(0) detected_names = [] detected_conf = [] # Process detections for i, det in enumerate(pred): # detections per image if len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(image.shape[2:], det[:, :4], src_image.shape).round() # Results for *xyxy, conf, cls in reversed(det): # Box w = float(xyxy[2] - xyxy[0]) h = float(xyxy[3] - xyxy[1]) prop_rect = core.GraphicsRectProperty() prop_rect.pen_color = self.colors[int(cls)] graphics_box = graphics_output.addRectangle(float(xyxy[0]), float(xyxy[1]), w, h, prop_rect) graphics_box.setCategory(self.names[int(cls)]) # Label name = self.names[int(cls)] prop_text = core.GraphicsTextProperty() prop_text.font_size = 8 prop_text.color = self.colors[int(cls)] graphics_output.addText(name, float(xyxy[0]), float(xyxy[1]), prop_text) detected_names.append(name) detected_conf.append(conf.item()) # Init numeric output numeric_ouput = self.getOutput(2) numeric_ouput.clearData() numeric_ouput.setOutputType(dataprocess.NumericOutputType.TABLE) numeric_ouput.addValueList(detected_conf, "Confidence", detected_names) self.emitStepProgress()
def run(self): self.beginTaskRun() # we use seed to keep the same color for our masks + boxes + labels (same random each time) random.seed(30) # Get input : input = self.getInput(0) srcImage = input.getImage() # Get output : output_image = self.getOutput(0) output_graph = self.getOutput(1) output_graph.setNewLayer("Detectron2_RetinaNet") # Get parameters : param = self.getParam() # predictor if not self.loaded: print("Chargement du modèle") if param.cuda == False: self.cfg.MODEL.DEVICE = "cpu" self.deviceFrom = "cpu" else: self.deviceFrom = "gpu" self.loaded = True self.predictor = DefaultPredictor(self.cfg) # reload model if CUDA check and load without CUDA elif self.deviceFrom == "cpu" and param.cuda == True: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(model_zoo.get_config_file( self.LINK_MODEL)) # load config from file(.yaml) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.LINK_MODEL) # download the model (.pkl) self.deviceFrom = "gpu" self.predictor = DefaultPredictor(self.cfg) # reload model if CUDA not check and load with CUDA elif self.deviceFrom == "gpu" and param.cuda == False: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.MODEL.DEVICE = "cpu" self.cfg.merge_from_file(model_zoo.get_config_file( self.LINK_MODEL)) # load config from file(.yaml) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.LINK_MODEL) # download the model (.pkl) self.deviceFrom = "cpu" self.predictor = DefaultPredictor(self.cfg) outputs = self.predictor(srcImage) # get outputs instances output_image.setImage(srcImage) boxes = outputs["instances"].pred_boxes scores = outputs["instances"].scores classes = outputs["instances"].pred_classes # to numpy if param.cuda: boxes_np = boxes.tensor.cpu().numpy() scores_np = scores.cpu().numpy() classes_np = classes.cpu().numpy() else: boxes_np = boxes.tensor.numpy() scores_np = scores.numpy() classes_np = classes.numpy() self.emitStepProgress() # keep only the results with proba > threshold scores_np_tresh = list() for s in scores_np: if float(s) > param.proba: scores_np_tresh.append(s) if len(scores_np_tresh) > 0: # text label with score labels = None class_names = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("thing_classes") if classes is not None and class_names is not None and len( class_names) > 1: labels = [class_names[i] for i in classes] if scores_np_tresh is not None: if labels is None: labels = [ "{:.0f}%".format(s * 100) for s in scores_np_tresh ] else: labels = [ "{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores_np_tresh) ] # Show Boxes + labels for i in range(len(scores_np_tresh)): color = [ random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255 ] prop_text = core.GraphicsTextProperty() prop_text.color = color prop_text.font_size = 7 output_graph.addText(labels[i], float(boxes_np[i][0]), float(boxes_np[i][1]), prop_text) prop_rect = core.GraphicsRectProperty() prop_rect.pen_color = color prop_rect.category = labels[i] output_graph.addRectangle( float(boxes_np[i][0]), float(boxes_np[i][1]), float(boxes_np[i][2] - boxes_np[i][0]), float(boxes_np[i][3] - boxes_np[i][1]), prop_rect) # Step progress bar: self.emitStepProgress() # Call endTaskRun to finalize process self.endTaskRun()
def run(self): self.beginTaskRun() # Get input : input = self.getInput(0) # Get output : output = self.getOutput(0) output_graph = self.getOutput(1) output_graph.setNewLayer("DensePose") srcImage = input.getImage() # Get parameters : param = self.getParam() # predictor if not self.loaded: print("Chargement du modèle") if param.cuda == False: self.cfg.MODEL.DEVICE = "cpu" self.deviceFrom = "cpu" else: self.deviceFrom = "gpu" self.loaded = True self.predictor = DefaultPredictor(self.cfg) # reload model if CUDA check and load without CUDA elif self.deviceFrom == "cpu" and param.cuda == True: print("Chargement du modèle") self.cfg = get_cfg() add_densepose_config(self.cfg) self.cfg.merge_from_file(self.folder + "/DensePose_git/configs/"+self.MODEL_NAME_CONFIG+".yaml") self.cfg.MODEL.WEIGHTS = self.folder + "/models/"+self.MODEL_NAME+".pkl" self.deviceFrom = "gpu" self.predictor = DefaultPredictor(self.cfg) # reload model if CUDA not check and load with CUDA elif self.deviceFrom == "gpu" and param.cuda == False: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.DEVICE = "cpu" add_densepose_config(self.cfg) self.cfg.merge_from_file(self.folder + "/DensePose_git/configs/"+self.MODEL_NAME_CONFIG+".yaml") self.cfg.MODEL.WEIGHTS = self.folder + "/models/"+self.MODEL_NAME+".pkl" self.deviceFrom = "cpu" self.predictor = DefaultPredictor(self.cfg) outputs = self.predictor(srcImage)["instances"] scores = outputs.get("scores").cpu() boxes_XYXY = outputs.get("pred_boxes").tensor.cpu() boxes_XYWH = BoxMode.convert(boxes_XYXY, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS) denseposes = outputs.get("pred_densepose").to_result(boxes_XYWH) # Number of iso values betwen 0 and 1 self.levels = np.linspace(0, 1, 9) cmap = cv2.COLORMAP_PARULA img_colors_bgr = cv2.applyColorMap((self.levels * 255).astype(np.uint8), cmap) self.level_colors_bgr = [ [int(v) for v in img_color_bgr.ravel()] for img_color_bgr in img_colors_bgr ] # text and rect graph properties properties_text = core.GraphicsTextProperty() properties_text.color = [255,255,255] properties_text.font_size = 10 properties_rect = core.GraphicsRectProperty() properties_rect.pen_color = [11,130,41] properties_line = core.GraphicsPolylineProperty() properties_line.line_size = 1 self.emitStepProgress() for i in range(len(denseposes)): if scores.numpy()[i] > param.proba: bbox_xywh = boxes_XYWH[i] bbox_xyxy = boxes_XYXY[i] result_encoded = denseposes.results[i] iuv_arr = DensePoseResult.decode_png_data(*result_encoded) # without indice surface self.visualize_iuv_arr(srcImage, iuv_arr, bbox_xywh, properties_line, output_graph) # with indice surface #self.visualize_iuv_arr_indiceSurface(srcImage, iuv_arr, bbox_xyxy, output_graph) output_graph.addRectangle(bbox_xyxy[0].item(), bbox_xyxy[1].item(), bbox_xyxy[2].item() - bbox_xyxy[0].item(), bbox_xyxy[3].item() - bbox_xyxy[1].item(),properties_rect) output_graph.addText(str(scores[i].item())[:5], float(bbox_xyxy[0].item()), float(bbox_xyxy[1].item()), properties_text) output.setImage(srcImage) self.emitStepProgress() self.endTaskRun()
def run(self): self.beginTaskRun() # we use seed to keep the same color for our boxes + labels (same random each time) random.seed(30) # Get input : input = self.getInput(0) srcImage = input.getImage() # Get output : output_graph = self.getOutput(1) output_graph.setNewLayer("KeypointRCNN") # Get parameters : param = self.getParam() # predictor if not self.loaded: print("Chargement du modèle") if param.cuda == False: self.cfg.MODEL.DEVICE = "cpu" self.deviceFrom = "cpu" else: self.deviceFrom = "gpu" self.predictor = DefaultPredictor(self.cfg) self.loaded = True # reload model if CUDA check and load without CUDA elif self.deviceFrom == "cpu" and param.cuda == True: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(model_zoo.get_config_file( self.LINK_MODEL)) # load config from file(.yaml) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.LINK_MODEL) # download the model (.pkl) self.predictor = DefaultPredictor(self.cfg) self.deviceFrom = "gpu" # reload model if CUDA not check and load with CUDA elif self.deviceFrom == "gpu" and param.cuda == False: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.DEVICE = "cpu" self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(model_zoo.get_config_file( self.LINK_MODEL)) # load config from file(.yaml) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.LINK_MODEL) # download the model (.pkl) self.predictor = DefaultPredictor(self.cfg) self.deviceFrom = "cpu" outputs = self.predictor(srcImage) # get outputs instances boxes = outputs["instances"].pred_boxes scores = outputs["instances"].scores classes = outputs["instances"].pred_classes keypoints = outputs["instances"].pred_keypoints # to numpy if param.cuda: boxes_np = boxes.tensor.cpu().numpy() scores_np = scores.cpu().numpy() classes_np = classes.cpu().numpy() keypoints_np = keypoints.cpu().numpy() else: boxes_np = boxes.tensor.numpy() scores_np = scores.numpy() classes_np = classes.numpy() keypoints_np = keypoints.numpy() self.emitStepProgress() # keep only the results with proba > threshold scores_np_tresh = list() for s in scores_np: if float(s) > param.proba: scores_np_tresh.append(s) if len(scores_np_tresh) > 0: # create random color for boxes and labels colors = [] for i in range(len(scores_np_tresh)): colors.append([ random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255 ]) # text label with score labels = None class_names = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("thing_classes") if classes is not None and class_names is not None and len( class_names) > 1: labels = [class_names[i] for i in classes] if scores_np_tresh is not None: if labels is None: labels = [ "{:.0f}%".format(s * 100) for s in scores_np_tresh ] else: labels = [ "{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores_np_tresh) ] # Show boxes + labels for i in range(len(scores_np_tresh)): properties_text = core.GraphicsTextProperty() properties_text.color = colors[ i] # start with i+1 we don't use the first color dedicated for the label mask properties_text.font_size = 7 properties_rect = core.GraphicsRectProperty() properties_rect.pen_color = colors[i] output_graph.addRectangle( float(boxes_np[i][0]), float(boxes_np[i][1]), float(boxes_np[i][2] - boxes_np[i][0]), float(boxes_np[i][3] - boxes_np[i][1]), properties_rect) output_graph.addText(labels[i], float(boxes_np[i][0]), float(boxes_np[i][1]), properties_text) self.emitStepProgress() # keypoints properties_point = core.GraphicsPointProperty() properties_point.pen_color = [0, 0, 0, 255] properties_point.brush_color = [0, 0, 255, 255] properties_point.size = 10 # get keypoints name if prob > Threshold keypoint_names = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("keypoint_names") for keypoints_obj in keypoints_np[:len(scores_np_tresh)]: visible_keypoints = {} for idx, kp in enumerate(keypoints_obj): x, y, prob = kp if prob > self._KEYPOINT_THRESHOLD: pts = core.CPointF(float(x), float(y)) output_graph.addPoint(pts, properties_point) if keypoint_names: keypoint_name = keypoint_names[idx] visible_keypoints[keypoint_name] = (x, y) # keypoints connections if MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]).get( "keypoint_connection_rules"): for kpName_0, kpName_1, color in MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get( "keypoint_connection_rules"): for kpName_0, kpName_1, color in MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get( "keypoint_connection_rules"): if kpName_0 in visible_keypoints and kpName_1 in visible_keypoints: x0, y0 = visible_keypoints[kpName_0] x1, y1 = visible_keypoints[kpName_1] color = [x for x in color] color.append(255) properties_line = core.GraphicsPolylineProperty( ) properties_line.pen_color = color pts0 = core.CPointF(float(x0), float(y0)) pts1 = core.CPointF(float(x1), float(y1)) lst_points = [pts0, pts1] output_graph.addPolyline( lst_points, properties_line) else: self.emitStepProgress() # Set input image to output 0 self.forwardInputImage(0, 0) # Step progress bar: self.emitStepProgress() # Call endTaskRun to finalize process self.endTaskRun()
def run(self): self.beginTaskRun() # we use seed to keep the same color for our masks + boxes + labels (same random each time) random.seed(30) # Get input : img_input = self.getInput(0) src_img = img_input.getImage() # Get output : mask_output = self.getOutput(0) output_graph = self.getOutput(2) output_graph.setImageIndex(1) output_graph.setNewLayer("MaskRCNN") # Get parameters : param = self.getParam() # predictor if not self.predictor or param.update_model: if param.dataset == "COCO": self.model_link = "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml" else: self.model_link = "Cityscapes/mask_rcnn_R_50_FPN.yaml" self.cfg = get_cfg() self.cfg.MODEL.DEVICE = param.device self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold # load config from file(.yaml) self.cfg.merge_from_file(model_zoo.get_config_file( self.model_link)) # download the model (.pkl) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.model_link) self.predictor = DefaultPredictor(self.cfg) param.update_model = False outputs = self.predictor(src_img) # get outputs instances boxes = outputs["instances"].pred_boxes scores = outputs["instances"].scores classes = outputs["instances"].pred_classes masks = outputs["instances"].pred_masks # to numpy boxes_np = boxes.tensor.cpu().numpy() scores_np = scores.cpu().numpy() # classes_np = classes.cpu().numpy() self.emitStepProgress() # keep only the results with proba > threshold scores_np_thresh = list() for s in scores_np: if float(s) > param.proba: scores_np_thresh.append(s) if len(scores_np_thresh) > 0: # create random color for masks + boxes + labels colors = [[0, 0, 0]] for i in range(len(scores_np_thresh)): colors.append([ random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255 ]) # text labels with scores labels = None class_names = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("thing_classes") if classes is not None and class_names is not None and len( class_names) > 1: labels = [class_names[i] for i in classes] if scores_np_thresh is not None: if labels is None: labels = [ "{:.0f}%".format(s * 100) for s in scores_np_thresh ] else: labels = [ "{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores_np_thresh) ] # Show boxes + labels for i in range(len(scores_np_thresh)): prop_text = core.GraphicsTextProperty() # start with i+1 we don't use the first color dedicated for the label mask prop_text.color = colors[i + 1] prop_text.font_size = 7 prop_rect = core.GraphicsRectProperty() prop_rect.pen_color = colors[i + 1] prop_rect.category = labels[i] output_graph.addRectangle( float(boxes_np[i][0]), float(boxes_np[i][1]), float(boxes_np[i][2] - boxes_np[i][0]), float(boxes_np[i][3] - boxes_np[i][1]), prop_rect) output_graph.addText(labels[i], float(boxes_np[i][0]), float(boxes_np[i][1]), prop_text) self.emitStepProgress() # label mask nb_objects = len(masks[:len(scores_np_thresh)]) if nb_objects > 0: masks = masks[:nb_objects, :, :, None] mask_or = masks[0] * nb_objects for j in range(1, nb_objects): mask_or = torch.max(mask_or, masks[j] * (nb_objects - j)) mask_numpy = mask_or.byte().cpu().numpy() mask_output.setImage(mask_numpy) # output mask apply to our original image # inverse colors to match boxes colors c = colors[1:] c = c[::-1] colors = [[0, 0, 0]] for col in c: colors.append(col) self.setOutputColorMap(1, 0, colors) else: self.emitStepProgress() self.forwardInputImage(0, 1) # Step progress bar: self.emitStepProgress() # Call endTaskRun to finalize process self.endTaskRun()
def run(self): self.beginTaskRun() # we use seed to keep the same color for our masks + boxes + labels (same random each time) random.seed(30) # Get input : input = self.getInput(0) srcImage = input.getImage() # Get output : output_graph = self.getOutput(2) output_graph.setImageIndex(1) output_graph.setNewLayer("PanopticSegmentation") # Get parameters : param = self.getParam() # predictor if not self.loaded: print("Chargement du modèle") if param.cuda == False: self.cfg.MODEL.DEVICE = "cpu" self.deviceFrom = "cpu" else: self.deviceFrom = "gpu" self.predictor = DefaultPredictor(self.cfg) self.loaded = True # reload model if CUDA check and load without CUDA elif self.deviceFrom == "cpu" and param.cuda == True: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(model_zoo.get_config_file( self.LINK_MODEL)) # load config from file(.yaml) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.LINK_MODEL) # download the model (.pkl) self.predictor = DefaultPredictor(self.cfg) self.deviceFrom = "gpu" # reload model if CUDA not check and load with CUDA elif self.deviceFrom == "gpu" and param.cuda == False: print("Chargement du modèle") self.cfg = get_cfg() self.cfg.MODEL.DEVICE = "cpu" self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(model_zoo.get_config_file( self.LINK_MODEL)) # load config from file(.yaml) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( self.LINK_MODEL) # download the model (.pkl) self.predictor = DefaultPredictor(self.cfg) self.deviceFrom = "cpu" outputs = self.predictor(srcImage)["panoptic_seg"] # get outputs of model mask = outputs[0] infos = outputs[1] # set mask output mask_output = self.getOutput(0) if param.cuda: mask_output.setImage(mask.cpu().numpy()) else: mask_output.setImage(mask.numpy()) self.emitStepProgress() # output visualisation nb_objects = len(infos) # create random color for masks + boxes + labels colors = [[0, 0, 0]] for i in range(nb_objects): colors.append([ random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255 ]) # get infos classes scores = list() classesThings = list() classesStuffs = list() labelsStuffs = list() for info in infos: if info["isthing"]: scores.append(info['score']) classesThings.append(info['category_id']) else: classesStuffs.append(info['category_id']) # text label with score - get classe name for thing and stuff from metedata labelsThings = None class_names = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("thing_classes") if classesThings is not None and class_names is not None and len( class_names) > 1: labelsThings = [class_names[i] for i in classesThings] if scores is not None: if labelsThings is None: labelsThings = ["{:.0f}%".format(s * 100) for s in scores] else: labelsThings = [ "{} {:.0f}%".format(l, s * 100) for l, s in zip(labelsThings, scores) ] class_names_stuff = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("stuff_classes") [labelsStuffs.append(class_names_stuff[x]) for x in classesStuffs] labels = labelsThings + labelsStuffs seg_ids = torch.unique(mask).tolist() self.emitStepProgress() # create masks - use for text_pos masks = list() for sid in seg_ids: if param.cuda: mymask = (mask == sid).cpu().numpy().astype(np.bool) else: mymask = (mask == sid).numpy().astype(np.bool) masks.append(mymask) # text pos = median of mask - median is less sensitive to outliers. if len(masks) > len( labels ): # unrecognized area - no given class for area labeled 0 for i in range(nb_objects): properties_text = core.GraphicsTextProperty() properties_text.color = colors[i + 1] properties_text.font_size = 7 text_pos = np.median(masks[i + 1].nonzero(), axis=1)[::-1] output_graph.addText(labels[i], text_pos[0], text_pos[1], properties_text) else: for i in range(nb_objects): properties_text = core.GraphicsTextProperty() properties_text.color = colors[i + 1] properties_text.font_size = 7 text_pos = np.median(masks[i].nonzero(), axis=1)[::-1] output_graph.addText(labels[i], text_pos[0], text_pos[1], properties_text) # output mask apply to our original image self.setOutputColorMap(1, 0, colors) self.forwardInputImage(0, 1) # Step progress bar: self.emitStepProgress() # Call endTaskRun to finalize process self.endTaskRun()
def run(self): self.beginTaskRun() # we use seed to keep the same color for our masks + boxes + labels (same random each time) random.seed(30) # Get input : input = self.getInput(0) srcImage = input.getImage() # Get output : mask_output = self.getOutput(0) output_graph = self.getOutput(2) output_graph.setImageIndex(1) output_graph.setNewLayer("PointRend") # Get parameters : param = self.getParam() # predictor if not self.loaded: print("Chargement du modèle") if param.cuda == False: self.cfg.MODEL.DEVICE = "cpu" self.deviceFrom = "cpu" else: self.deviceFrom = "gpu" self.loaded = True self.predictor = DefaultPredictor(self.cfg) # reload model if CUDA check and load without CUDA elif self.deviceFrom == "cpu" and param.cuda == True: print("Chargement du modèle") self.cfg = get_cfg() add_pointrend_config(self.cfg) self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(self.folder + self.path_to_config) self.cfg.MODEL.WEIGHTS = self.folder + self.path_to_model self.deviceFrom = "gpu" self.predictor = DefaultPredictor(self.cfg) # reload model if CUDA not check and load with CUDA elif self.deviceFrom == "gpu" and param.cuda == False: print("Chargement du modèle") self.cfg = get_cfg() add_pointrend_config(self.cfg) self.cfg.MODEL.DEVICE = "cpu" self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.threshold self.cfg.merge_from_file(self.folder + self.path_to_config) self.cfg.MODEL.WEIGHTS = self.folder + self.path_to_model self.deviceFrom = "cpu" self.predictor = DefaultPredictor(self.cfg) outputs = self.predictor(srcImage) # get outputs instances boxes = outputs["instances"].pred_boxes scores = outputs["instances"].scores classes = outputs["instances"].pred_classes masks = outputs["instances"].pred_masks # to numpy if param.cuda: boxes_np = boxes.tensor.cpu().numpy() scores_np = scores.cpu().numpy() classes_np = classes.cpu().numpy() else: boxes_np = boxes.tensor.numpy() scores_np = scores.numpy() classes_np = classes.numpy() self.emitStepProgress() # keep only the results with proba > threshold scores_np_tresh = list() for s in scores_np: if float(s) > param.proba: scores_np_tresh.append(s) if len(scores_np_tresh) > 0: # create random color for masks + boxes + labels colors = [[0, 0, 0]] for i in range(len(scores_np_tresh)): colors.append([ random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255 ]) # text labels with scores labels = None class_names = MetadataCatalog.get( self.cfg.DATASETS.TRAIN[0]).get("thing_classes") if classes is not None and class_names is not None and len( class_names) > 1: labels = [class_names[i] for i in classes] if scores_np_tresh is not None: if labels is None: labels = [ "{:.0f}%".format(s * 100) for s in scores_np_tresh ] else: labels = [ "{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores_np_tresh) ] # Show boxes + labels for i in range(len(scores_np_tresh)): prop_text = core.GraphicsTextProperty() # start with i+1 we don't use the first color dedicated for the label mask prop_text.color = colors[i + 1] prop_text.font_size = 7 prop_rect = core.GraphicsRectProperty() prop_rect.pen_color = colors[i + 1] prop_rect.category = labels[i] output_graph.addRectangle( float(boxes_np[i][0]), float(boxes_np[i][1]), float(boxes_np[i][2] - boxes_np[i][0]), float(boxes_np[i][3] - boxes_np[i][1]), prop_rect) output_graph.addText(labels[i], float(boxes_np[i][0]), float(boxes_np[i][1]), prop_text) self.emitStepProgress() # label mask nb_objects = len(masks[:len(scores_np_tresh)]) if nb_objects > 0: masks = masks[:nb_objects, :, :, None] mask_or = masks[0] * nb_objects for j in range(1, nb_objects): mask_or = torch.max(mask_or, masks[j] * (nb_objects - j)) mask_numpy = mask_or.byte().cpu().numpy() mask_output.setImage(mask_numpy) # output mask apply to our original image # inverse colors to match boxes colors c = colors[1:] c = c[::-1] colors = [[0, 0, 0]] for col in c: colors.append(col) self.setOutputColorMap(1, 0, colors) else: self.emitStepProgress() self.forwardInputImage(0, 1) # Step progress bar: self.emitStepProgress() # Call endTaskRun to finalize process self.endTaskRun()
def run(self): # Core function of your process # Call beginTaskRun for initialization self.beginTaskRun() random.seed(1) # Get parameters : param = self.getParam() # Get input : input = self.getInput(0) src_image = input.getImage() h = src_image.shape[0] w = src_image.shape[1] # Step progress bar: self.emitStepProgress() # Load model if self.model is None or param.update: # Load class names self.load_class_names() # Load model use_torchvision = param.dataset != "Custom" self.model = models.mask_rcnn(use_pretrained=use_torchvision, classes=len(self.class_names)) if param.dataset == "Custom": self.model.load_state_dict(torch.load(param.model_path)) self.model.to(self.device) param.update = False pred = self.predict(src_image) cpu = torch.device("cpu") boxes = pred[0]["boxes"].to(cpu).numpy().tolist() scores = pred[0]["scores"].to(cpu).numpy().tolist() labels = pred[0]["labels"].to(cpu).numpy().tolist() masks = pred[0]["masks"] # Step progress bar: self.emitStepProgress() # Forward input image to result image self.forwardInputImage(0, 1) # Init graphics output graphics_output = self.getOutput(2) graphics_output.setNewLayer("MaskRCNN") graphics_output.setImageIndex(1) prop_text = core.GraphicsTextProperty() prop_rect = core.GraphicsRectProperty() # Get predictions detected_names = [] detected_scores = [] size = masks.size() valid_results = [scores.index(x) for x in scores if x > param.confidence] object_value = len(valid_results) mask_or = torch.zeros(1, size[2], size[3]).to(device=self.device) colors = [[0, 0, 0]] for i in valid_results: # color color = [random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] prop_text.color = color prop_rect.pen_color = color # box box = boxes[i] w = box[2] - box[0] h = box[3] - box[1] graphics_box = graphics_output.addRectangle(float(box[0]), float(box[1]), float(w), float(h), prop_rect) graphics_box.setCategory(self.class_names[labels[i]]) # label label = self.class_names[labels[i]] + ": {:.3f}".format(scores[i]) graphics_output.addText(label, float(box[0]), float(box[1]), prop_text) detected_names.append(self.class_names[labels[i]]) detected_scores.append(scores[i]) # masks -> merge into a single labelled image mask = (masks[i] > param.mask_threshold).float() mask_or = torch.max(mask_or, mask * object_value) object_value -= 1 colors.insert(1, color) # Segmentation mask output mask_output = self.getOutput(0) mask_numpy = mask_or.squeeze().byte().cpu().numpy() mask_output.setImage(mask_numpy) self.setOutputColorMap(1, 0, colors) # Init numeric output numeric_ouput = self.getOutput(3) numeric_ouput.clearData() numeric_ouput.setOutputType(dataprocess.NumericOutputType.TABLE) numeric_ouput.addValueList(detected_scores, "Probability", detected_names) # Step progress bar: self.emitStepProgress() # Call endTaskRun to finalize process self.endTaskRun()
def manage_outputs(predictions, img, param, graphics_output): crop_mask = True display_masks = True display_text = True display_bboxes = True display_scores = True # Put values in range [0 - 1] h, w, _ = img.shape # Post-processing save = cfg.rescore_bbox cfg.rescore_bbox = True t = postprocess(predictions, w, h, visualize_lincomb=False, crop_masks=crop_mask, score_threshold=param.confidence) cfg.rescore_bbox = save # Copy idx = t[1].argsort(0, descending=True)[:param.top_k] if cfg.eval_mask_branch: # Masks are drawn on the GPU, so don't copy masks = t[3][idx] classes, scores, boxes = [x[idx].cpu().numpy() for x in t[:3]] # Filter available detections num_dets_to_consider = min(param.top_k, classes.shape[0]) for j in range(num_dets_to_consider): if scores[j] < param.confidence: num_dets_to_consider = j break # Quick and dirty lambda for selecting the color for a particular index # Also keeps track of a per-gpu color cache for maximum speed class_color = False def get_color(j, on_gpu=None): global color_cache color_idx = (classes[j] * 5 if class_color else j * 5) % len(COLORS) if on_gpu is not None and color_idx in color_cache[on_gpu]: return color_cache[on_gpu][color_idx] else: color = COLORS[color_idx] # The image might come in as RGB or BRG, depending color = (color[2], color[1], color[0]) if on_gpu is not None: color = torch.Tensor(color).to(on_gpu).float() / 255. color_cache[on_gpu][color_idx] = color return color # First, draw the masks on the GPU where we can do it really fast # Beware: very fast but possibly unintelligible mask-drawing code ahead # I wish I had access to OpenGL or Vulkan but alas, I guess Pytorch tensor operations will have to suffice if display_masks and cfg.eval_mask_branch and num_dets_to_consider > 0: # After this, mask is of size [num_dets, h, w, 1] masks = masks[:num_dets_to_consider, :, :, None] # Cumulative mask # For each object, we create a mask with label values for display purpose # For overlapping objects, we take those with the better probability with a MAX operator # that's why we reverse the mask weights mask_or = masks[0] * num_dets_to_consider for j in range(1, num_dets_to_consider): mask_or = torch.max(mask_or, masks[j] * (num_dets_to_consider - j)) # Get the numpy array of the mask mask_numpy = mask_or.byte().cpu().numpy() colorvec = [[0, 0, 0]] if num_dets_to_consider == 0: return mask_numpy, colorvec if display_text or display_bboxes: for j in reversed(range(num_dets_to_consider)): x1, y1, x2, y2 = boxes[j, :] color = get_color(j) colorvec.append(list(color)) score = scores[j] if display_bboxes: rect_prop = core.GraphicsRectProperty() rect_prop.pen_color = list(color) graphics_box = graphics_output.addRectangle( float(x1), float(y1), float(x2 - x1), float(y2 - y1), rect_prop) graphics_box.setCategory(cfg.dataset.class_names[classes[j]]) if display_text: _class = cfg.dataset.class_names[classes[j]] text_str = '%s: %.2f' % (_class, score) if display_scores else _class text_prop = core.GraphicsTextProperty() text_prop.bold = True graphics_output.addText(text_str, float(x1), float(y1), text_prop) return mask_numpy, colorvec