def predict_proba(self, data_dir, data_names=None, transform=None, batch_size=1, conf_thres=0.7, nms_thres=0.4, verbose=True, **kwargs): """ During inference, the model requires only the input tensors, and returns the post-processed predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as follows: - boxes (Tensor[N, 4]): the predicted boxes in [x0, y0, x1, y1] format, with values between 0 and H and 0 and W - labels (Tensor[N]): the predicted labels for each image - scores (Tensor[N]): the scores or each prediction """ register_all_cityscapes(data_dir) self.cfg.DATASETS.TEST = ["cityscapes_fine2_instance_seg_test"] self.cfg.MODEL.WEIGHTS = os.path.join(self.cfg.OUTPUT_DIR, 'model_final.pth') self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST =conf_thres predictor = DefaultPredictor(self.cfg) data_loader = LiuyTrainer.build_test_loader(self.cfg, "cityscapes_fine2_instance_seg_test") results = [] for batch in data_loader: for item in batch: file_name = item['file_name'] img = cv2.imread(file_name) prediction = predictor(img) record = {'boxes': prediction['instances'].pred_boxes, 'labels': prediction['instances'].pred_classes, \ 'scores': prediction['instances'].scores} results.append(record) return results
def __init__(self, args, project_id, model_name=None, num_classes=None, pytorch_model=None): self.args = args self.project_id = project_id self.model_name = model_name self.num_classes =num_classes self.data_dir = None self.lr = None # prepare cfg for build model self.cfg = setup(args=args, project_id=project_id, model_name=model_name, num_classes=num_classes) super(Detctron2AlObjDetModel, self).__init__(project_id) self.model, self.device = load_prj_model(project_id=project_id) if self.model is None: if pytorch_model: assert isinstance( pytorch_model, nn.Module), 'pytorch_model must inherit from torch.nn.Module' self.model = pytorch_model print("get a pre-trained model from parameter for project{}".format(project_id)) else: assert model_name in MODEL_NAME.keys( ), 'model_name must be one of {}'.format(MODEL_NAME.keys()) if not num_classes: raise ValueError( "Deep model of project {} is not initialized, please specify the model name and number of classes.".format( project_id)) self.model = LiuyTrainer.build_model(self.cfg) self.model = self.model.to(self.device) print("Initialize a pre-trained model for project{}".format(project_id)) else: print("load project {} model from file".format(project_id)) print(self.model)
def test(self): self.cfg.DATASETS.TEST = ["cityscapes_fine2_instance_seg_val"] self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.test(self.cfg, self.trainer.model)
def fit_on_subset(self, **kwargs): self.cfg.DATASETS.TRAIN = ["cityscapes_fine2_instance_seg_sub_train"] self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.resume_or_load(resume=args.resume) self.trainer.train() self.save_model()
class InsSegModel(BaseInsSegModel): """Mask_RCNN""" def __init__(self, args, project_id, data_dir): self.args = args super(InsSegModel, self).__init__(project_id, data_dir) # tow ways to get model # 1:load the model which has been trained # 2:use the function:LiuyTrainer.build_model(self.cfg) self.model, self.device = load_prj_model(project_id=project_id) self.cfg = setup(args=args, project_id=project_id, data_dir=data_dir) if self.model is None: self.model = LiuyTrainer.build_model(self.cfg) self.model = self.model.to(self.device) print("Initialize a pre-trained model for project{}".format( project_id)) else: print("load project {} model from file".format(project_id)) self.trainer = LiuyTrainer(self.cfg, self.model) def fit(self): self.cfg.DATASETS.TRAIN = ["cityscapes_fine2_instance_seg_train"] self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.resume_or_load(resume=args.resume) self.trainer.train() self.save_model() def fit_on_subset(self, **kwargs): self.cfg.DATASETS.TRAIN = ["cityscapes_fine2_instance_seg_sub_train"] self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.resume_or_load(resume=args.resume) self.trainer.train() self.save_model() def test(self): self.cfg.DATASETS.TEST = ["cityscapes_fine2_instance_seg_val"] self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.test(self.cfg, self.trainer.model) def compute_loss(self, image_dir, gt_dir): register_a_cityscapes(image_dir, gt_dir, 'dataset_name') self.cfg.DATASETS.TRAIN = ["dataset_name"] computer = LiuyComputeLoss(self.cfg, self.model) return computer.compute() def predict_proba(self, image_dir, gt_dir, conf_thres=0.7, nms_thres=0.4, verbose=True, **kwargs): """ During inference, the model requires only the input tensors, and returns the post-processed predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as follows: - boxes (Tensor[N, 4]): the predicted boxes in [x0, y0, x1, y1] format, with values between 0 and H and 0 and W - labels (Tensor[N]): the predicted labels for each image - scores (Tensor[N]): the scores or each prediction """ self.cfg.MODEL.WEIGHTS = os.path.join(self.cfg.OUTPUT_DIR, 'model_final.pth') self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = conf_thres register_a_cityscapes(image_dir, gt_dir, 'dataset_name') predictor = DefaultPredictor(self.cfg) data_loader = LiuyTrainer.build_test_loader(self.cfg, "dataset_name") results = [] for batch in data_loader: for item in batch: file_name = item['file_name'] img = cv2.imread(file_name) prediction = predictor(img) record = {'file_name': file_name, 'boxes': prediction['instances'].pred_boxes, 'labels': prediction['instances'].pred_classes, \ 'scores': prediction['instances'].scores} results.append(record) return results def predict(self, image_dir, gt_dir): '''predict :param data_dir: str The path to the data folder. :param data_names: list, optional (default=None) The data names. If not specified, it will all the files in the data_dir. :param transform: torchvision.transforms.Compose, optional (default=None) Transforms object that will be applied to the image data. :return: pred: 1D array The prediction result. Shape [n_samples] ''' proba_result = self.predict_proba(image_dir, gt_dir) return proba_result def save_model(self): with open( os.path.join(TRAINED_MODEL_DIR, self._proj_id + '_model.pkl'), 'wb') as f: pickle.dump(self.trainer.model, f)
class Detctron2AlObjDetModel(BaseDeepModel): """Faster_RCNN""" def __init__(self, args, project_id, model_name=None, num_classes=None, pytorch_model=None): self.args = args self.project_id = project_id self.model_name = model_name self.num_classes = num_classes self.data_dir = None self.lr = None # prepare cfg for build model self.cfg = setup(args=args, project_id=project_id, model_name=model_name, num_classes=num_classes) super(Detctron2AlObjDetModel, self).__init__(project_id) self.model, self.device = load_prj_model(project_id=project_id) if self.model is None: if pytorch_model: assert isinstance( pytorch_model, nn.Module ), 'pytorch_model must inherit from torch.nn.Module' self.model = pytorch_model print("get a pre-trained model from parameter for project{}". format(project_id)) else: assert model_name in MODEL_NAME.keys( ), 'model_name must be one of {}'.format(MODEL_NAME.keys()) if not num_classes: raise ValueError( "Deep model of project {} is not initialized, please specify the model name and number of classes." .format(project_id)) self.model = LiuyTrainer.build_model(self.cfg) self.model = self.model.to(self.device) print("Initialize a pre-trained model for project{}".format( project_id)) else: print("load project {} model from file".format(project_id)) print(self.model) def fit(self, data_dir, label=None, transform=None, batch_size=1, shuffle=False, data_names=None, optimize_method='Adam', optimize_param=None, loss='CrossEntropyLoss', loss_params=None, num_epochs=10, save_model=True, test_label=None, **kwargs): self.data_dir = data_dir print("Command Line Args:", args) self.cfg = setup(args, project_id=self.project_id, model_name=self.model_name, num_classes=self.num_classes, data_dir=data_dir) self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.resume_or_load(resume=args.resume) self.trainer.train() self.save_model() def fit2(self, data_dir, label=None, transform=None, batch_size=1, shuffle=False, data_names=None, optimize_method='Adam', optimize_param=None, loss='CrossEntropyLoss', loss_params=None, num_epochs=10, save_model=True, test_label=None, **kwargs): self.data_dir = data_dir print("Command Line Args:", args) self.cfg = setup(args, project_id=self.project_id, model_name=self.model_name, num_classes=self.num_classes, data_dir=data_dir) self.trainer = LiuyTrainer(self.cfg, self.model) self.trainer.resume_or_load(resume=args.resume) self.save_model() def predict_proba(self, data_dir, data_names=None, transform=None, batch_size=1, conf_thres=0.7, nms_thres=0.4, verbose=True, **kwargs): """ During inference, the model requires only the input tensors, and returns the post-processed predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as follows: - boxes (Tensor[N, 4]): the predicted boxes in [x0, y0, x1, y1] format, with values between 0 and H and 0 and W - labels (Tensor[N]): the predicted labels for each image - scores (Tensor[N]): the scores or each prediction """ #测试权重 # self.cfg.MODEL.WEIGHTS = os.path.join('/media/tangyp/Data/model_file/OUTPUT_DIR', 'model_final.pth') #实际运行权重 self.cfg.MODEL.WEIGHTS = os.path.join(self.cfg.OUTPUT_DIR, 'model_final.pth') self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = conf_thres predictor = DefaultPredictor(self.cfg) DatasetCatalog.register( "custom_val", lambda data_dir=data_dir: get_custom_dicts(data_dir)) data_loader = LiuyTrainer.build_test_loader(self.cfg, "custom_val") results = [] for batch in data_loader: for item in batch: file_name = item['file_name'] img = cv2.imread(file_name) prediction = predictor(img) record = {'boxes': prediction['instances'].pred_boxes, 'labels': prediction['instances'].pred_classes, \ 'scores': prediction['instances'].scores} results.append(record) return results def predict(self, data_dir, data_names=None, transform=None): '''predict :param data_dir: str The path to the data folder. :param data_names: list, optional (default=None) The data names. If not specified, it will all the files in the data_dir. :param transform: torchvision.transforms.Compose, optional (default=None) Transforms object that will be applied to the image data. :return: pred: 1D array The prediction result. Shape [n_samples] ''' proba_result = self.predict_proba(data_dir=data_dir, data_names=data_names, transform=transform) return proba_result def test(self, data_dir, label, batch_size, **kwargs): self.model_ft.eval() assert isinstance(label, dict) dataloader = create_faster_rcnn_dataloader(data_dir=data_dir, label_dict=label, augment=False, batch_size=batch_size, shuffle=False) with torch.no_grad(): return evaluate(self.model_ft, dataloader, self.device) def save_model(self): with open( os.path.join(TRAINED_MODEL_DIR, self._proj_id + '_model.pkl'), 'wb') as f: pickle.dump(self.trainer.model, f)