def __init__(
self,
config_file: Optional[str] = None,
weights: Optional[str] = None,
threshold: float = 0.5,
image_library="PIL",
):
cfg = get_cfg()
if config_file == "model_zoo/coco-detection":
cfg.merge_from_file(get_config_file(COCO_CFG))
cfg.MODEL.WEIGHTS = get_checkpoint_url(COCO_CFG)
if weights is not None:
logger.warning(
"Weights should not be provided when selecting a model from the zoo."
)
elif config_file == "model_zoo/lvis-segmentation":
cfg.merge_from_file(get_config_file(LVIS_CFG))
cfg.MODEL.WEIGHTS = get_checkpoint_url(LVIS_CFG)
if weights is not None:
logger.warning(
"Weights should not be provided when selecting a model from the zoo."
)
else:
cfg.merge_from_file(config_file)
cfg.MODEL.WEIGHTS = weights
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
logger.info(f"Detectron2 configuration:\n{cfg}")
self._patch_detectron()
self.d2 = DefaultPredictor(cfg)
self.image_library = image_library
def __init__(self, model_type="IS"):
self.cfg = get_cfg()
self.model_type = model_type
# load model config and pretrained model
if model_type == "IS": # instance segmentation
self.cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
elif model_type == "KP": # keypoint detection
self.cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
elif model_type == "PS": # panoptics segmentation
self.cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml"))
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml")
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
self.cfg.MODEL.DEVICE = "cpu" # cpu or cuda
self.predictor = DefaultPredictor(self.cfg)
def train(self, training_frames, network, get_images):
if (training_frames <= 0):
raise ValueError(
"The number of input frames must be bigger than 0")
self.cfg.OUTPUT_DIR = (f'../datasets/detectron2/{network}')
self.generate_datasets(training_frames, get_images)
retinanet_path = "COCO-Detection/retinanet_R_101_FPN_3x.yaml"
faster_rcnn_path = "COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"
if (detectron2.__version__ == "0.1"):
if network == 'faster_rcnn':
self.cfg.merge_from_file(
pkg_resources.resource_filename(
"detectron2.model_zoo",
os.path.join("configs", faster_rcnn_path)))
self.cfg.MODEL.WEIGHTS = model_zoo.ModelZooUrls.get(
faster_rcnn_path)
if network == 'retinanet':
self.cfg.merge_from_file(
pkg_resources.resource_filename(
"detectron2.model_zoo",
os.path.join("configs", retinanet_path)))
self.cfg.MODEL.WEIGHTS = model_zoo.ModelZooUrls.get(
retinanet_path)
else:
if network == 'faster_rcnn':
self.cfg.merge_from_file(
model_zoo.get_config_file(faster_rcnn_path))
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
faster_rcnn_path)
if network == 'retinanet':
self.cfg.merge_from_file(
model_zoo.get_config_file(retinanet_path))
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
retinanet_path)
self.cfg.DATASETS.TRAIN = ('train_set', )
self.cfg.DATASETS.TEST = ('val_set', )
self.cfg.DATALOADER.NUM_WORKERS = 1
self.cfg.SOLVER.IMS_PER_BATCH = 1
self.cfg.SOLVER.BASE_LR = 0.001
self.cfg.SOLVER.MAX_ITER = 2000
self.cfg.SOLVER.STEPS = (1000, 2000)
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
self.cfg.MODEL.DEVICE = 'cuda'
if not os.path.isfile(
os.path.join(self.cfg.OUTPUT_DIR, 'model_final.pth')):
os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(self.cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def __initialize_network(self, network, gt_frames):
if self.training:
self.train(self.thr_n_of_training, self.method_train, network,
gt_frames)
retinanet_path = "COCO-Detection/retinanet_R_101_FPN_3x.yaml"
faster_rcnn_path = "COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"
if (detectron2.__version__ == "0.1"):
if network == "retinanet":
self.cfg.merge_from_file(
pkg_resources.resource_filename(
"detectron2.model_zoo",
os.path.join("configs", retinanet_path)))
if not self.training:
self.cfg.MODEL.WEIGHTS = model_zoo.ModelZooUrls.get(
retinanet_path)
else:
self.cfg.MODEL.WEIGHTS = os.path.join(
self.cfg.OUTPUT_DIR, 'model_final.pth')
self.cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.5 # set threshold for this model
if network == "faster_rcnn":
self.cfg.merge_from_file(
pkg_resources.resource_filename(
"detectron2.model_zoo",
os.path.join("configs", faster_rcnn_path)))
if not self.training:
self.cfg.MODEL.WEIGHTS = model_zoo.ModelZooUrls.get(
faster_rcnn_path)
else:
self.cfg.MODEL.WEIGHTS = os.path.join(
self.cfg.OUTPUT_DIR, 'model_final.pth')
self.cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.5 # set threshold for this model
else:
if network == "retinanet":
self.cfg.merge_from_file(
model_zoo.get_config_file(retinanet_path))
if not self.training:
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
retinanet_path)
else:
self.cfg.MODEL.WEIGHTS = os.path.join(
self.cfg.OUTPUT_DIR, 'model_final.pth')
self.cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.5 # set threshold for this model
if network == "faster_rcnn":
self.cfg.merge_from_file(
model_zoo.get_config_file(faster_rcnn_path))
if not self.training:
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
faster_rcnn_path)
else:
self.cfg.MODEL.WEIGHTS = os.path.join(
self.cfg.OUTPUT_DIR, 'model_final.pth')
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Create predictor
return DefaultPredictor(self.cfg)
def setWeights(args, cfg):
if args.weights.lower() == "last":
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
else:
if args.type.lower() == "maskrcnn":
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
else:
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
def main(args):
cfg = setup(args)
dataset_root_dir = args.dataset_dir
dataset_name = dataset_root_dir.split('/')[-1]
image_dir = os.path.join(dataset_root_dir, "data_semantics/")
gt_dir = os.path.join(dataset_root_dir, "gtFine/")
# Register Dataset
splits = ["train", "val"] if args.do_eval else ["train"]
# Set from_json=False because there is no gt json for KITTI
register_dataset_instance(image_dir,
gt_dir,
splits=splits,
dataset_name=dataset_name,
from_json=False)
dataset_train = dataset_name + "_instance_train"
dataset_val = dataset_name + "_instance_val"
# Specify and create output directory
cfg.OUTPUT_DIR = "{}/output_{}".format(args.output_dir, args.backbone)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
# Use the registered dataset for training and validation
cfg.DATASETS.TRAIN = (dataset_train, )
cfg.DATASETS.TEST = (dataset_val, ) if args.do_eval else ()
cfg.TEST.EVAL_PERIOD = 600
cfg.DATALOADER.NUM_WORKERS = 4
if args.ckpt_path is None:
# load pre-trained weights
if args.backbone == "resnet-50":
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"Cityscapes/mask_rcnn_R_50_FPN.yaml")
elif args.backbone == "resnext-152":
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml")
else:
cfg.MODEL.WEIGHTS = args.ckpt_path
# config training parameters
cfg.SOLVER.IMS_PER_BATCH = 16
cfg.SOLVER.BASE_LR = 0.01
cfg.SOLVER.GAMMA = 0.1
cfg.SOLVER.STEPS = (
10000, 20000
) # iteration numbers to decrease learning rate by SOLVER.GAMMA
cfg.SOLVER.MAX_ITER = 30000
# choose trainer
trainer = MyTrainer(cfg)
trainer.resume_or_load(resume=False)
return trainer.train()
def _get_config(self):
cfg = get_cfg()
cfg.set_new_allowed(True)
# augment and initialize Detectron2 cfg with selected model
defaultConfig = optionsHelper.get_hierarchical_value(
self.options, ['options', 'model', 'config', 'value', 'id'])
if isinstance(defaultConfig, str):
# try to load from Detectron2's model zoo
try:
configFile = model_zoo.get_config_file(defaultConfig)
except:
# not available; try to load locally instead
configFile = os.path.join(
os.getcwd(), 'ai/models/detectron2/_functional/configs',
defaultConfig)
if not os.path.exists(configFile):
configFile = None
if configFile is not None:
cfg.merge_from_file(configFile)
try:
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(defaultConfig)
except:
pass
return cfg
def __init__(self,
dataset_name: str = None,
coco_json_path: str = None,
image_root: str = None,
is_predictor: bool = False):
self._dataset_name = dataset_name
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
) # Let training initialize from model zoo
if is_predictor:
self.cfg = cfg
self.set_predictor(is_original=True)
else:
if self._dataset_name is None:
raise Exception("dataset name is None !")
cfg.DATASETS.TRAIN = (dataset_name, )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.01 # pick a good LR
cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 10 # only has one class (ballon)
register_coco_instances(dataset_name, {}, coco_json_path,
image_root)
super().__init__(self.cfg)
self.resume_or_load(resume=False)
def initialization():
"""Loads configuration and model for the prediction.
Returns:
cfg (detectron2.config.config.CfgNode): Configuration for the model.
predictor (detectron2.engine.defaults.DefaultPredicto): Model to use.
by the model.
"""
cfg = get_cfg()
# Force model to operate within CPU, erase if CUDA compatible devices ara available
cfg.MODEL.DEVICE = 'cpu'
# Add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
# Set threshold for this model
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
# Initialize prediction model
predictor = DefaultPredictor(cfg)
return cfg, predictor
def setup_cfg(config_file,
weights_file=None,
config_opts=[],
confidence_threshold=None,
cpu=False):
# load config from file and command-line arguments
cfg = get_cfg()
cfg.merge_from_file(config_file)
cfg.merge_from_list(config_opts)
if confidence_threshold is not None:
# Set score_threshold for builtin models
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = confidence_threshold
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = confidence_threshold
if weights_file is not None:
cfg.MODEL.WEIGHTS = weights_file
else:
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
config_file.replace("configs/", ""))
if cpu or not torch.cuda.is_available():
cfg.MODEL.DEVICE = "cpu"
cfg.freeze()
return cfg
def load_model(
model_url: str = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml",
threshold: float = 0.7,
return_cfg: bool = False) -> DefaultPredictor:
"""Loads detectron2 model at model dir with threshold, option to return
:param model_url: Points to yaml file containing config, defaults to "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
:type model_url: str, optional
:param threshold: Confidence threshold, defaults to 0.7
:type threshold: float, optional
:param return_cfg: If CfgNode obj should be returned, defaults to False
:type return_cfg: bool, optional
:return: Detectron2 default predictor
:rtype: DefaultPredictor
"""
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(model_url))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_url)
try:
predictor = DefaultPredictor(cfg)
except AssertionError as a:
if "NVIDIA driver" in str(a):
cfg.MODEL.DEVICE = "cpu"
predictor = DefaultPredictor(cfg)
else:
raise a
if return_cfg:
return predictor, cfg
else:
return predictor
def setup_cfg(config_file: str, confidence_threshold=0.5):
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(config_file))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_file)
cfg.freeze()
return cfg
def train_model(trainset_name: str, learning_rate: int, num_iteration: int,
batch_per_image: int, num_classes: int):
cfg = get_cfg()
cfg.MODEL.DEVICE = "cpu"
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = (trainset_name, )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = learning_rate # pick a good LR
cfg.SOLVER.MAX_ITER = num_iteration # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = batch_per_image # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes # (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
# Inference should use the config with parameters that are used in training
# cfg now already contains everything we've set previously. We changed it a little bit for inference:
cfg.MODEL.WEIGHTS = os.path.join(
cfg.OUTPUT_DIR, "model_final.pth") # path to the model we just trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold
predictor = DefaultPredictor(cfg)
return predictor
def test_get_url(self):
url = model_zoo.get_checkpoint_url(
"Misc/scratch_mask_rcnn_R_50_FPN_3x_gn.yaml")
self.assertEqual(
url,
"https://dl.fbaipublicfiles.com/detectron2/Misc/scratch_mask_rcnn_R_50_FPN_3x_gn/138602908/model_final_01ca85.pkl", # noqa
)
def task_a(model_name, model_file, checkpoint=None, evaluate=True, visualize=True):
print('Running task A for model', model_name)
if checkpoint:
SAVE_PATH = os.path.join('./results_week_5_task_a', model_name + '_wCheckpoint')
else:
SAVE_PATH = os.path.join('./results_week_5_task_a', model_name)
os.makedirs(SAVE_PATH, exist_ok=True)
# Loading data
print('Loading data')
dataloader = MOTS_Dataloader(dataset='motschallenge')
def mots_train(): return dataloader.get_dicts(train_flag=True)
def mots_val(): return dataloader.get_dicts(train_flag=False)
DatasetCatalog.register('MOTS_train', mots_train)
MetadataCatalog.get('MOTS_train').set(thing_classes=list(MOTS_CATEGORIES.keys()))
DatasetCatalog.register('MOTS_val', mots_val)
MetadataCatalog.get('MOTS_val').set(thing_classes=list(MOTS_CATEGORIES.keys()))
# Load model and configuration
print('Loading Model')
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(model_file))
model_training_metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]) # Store current model training metadata
cfg.DATASETS.TRAIN = ('MOTS_train', )
cfg.DATASETS.TEST = ('MOTS_val', )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.OUTPUT_DIR = SAVE_PATH
if checkpoint:
print('Using Checkpoint')
cfg.MODEL.WEIGHTS = checkpoint
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
else:
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_file)
if evaluate:
model = build_model(cfg)
DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)
# Evaluation
print('Evaluating')
evaluator = COCOEvaluator('MOTS_val', cfg, False, output_dir=SAVE_PATH)
trainer = DefaultTrainer(cfg)
trainer.test(cfg, model, evaluators=[evaluator])
if visualize:
# Qualitative results: visualize some results
print('Getting qualitative results')
predictor = DefaultPredictor(cfg)
inputs = mots_val()
inputs = inputs[:20] + inputs[-20:]
for i, input in enumerate(inputs):
img = cv2.imread(input['file_name'])
outputs = predictor(img)
v = Visualizer(
img[:, :, ::-1],
metadata=model_training_metadata,
scale=0.8,
instance_mode=ColorMode.IMAGE)
v = v.draw_instance_predictions(outputs['instances'].to('cpu'))
cv2.imwrite(os.path.join(SAVE_PATH, 'Inference_' + model_name + '_inf_' + str(i) + '.png'), v.get_image()[:, :, ::-1])
def do_training(train):
cfg = get_cfg()
# cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
# cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_101_FPN_3x.yaml"))
# cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x.yaml"))
# cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.merge_from_file(model_zoo.get_config_file("Misc/scratch_mask_rcnn_R_50_FPN_3x_gn.yaml"))
cfg.DATASETS.TRAIN = ("cuboid_dataset_train",)
cfg.DATASETS.TEST = ("cuboid_dataset_val",) #used by evaluator- do Not remove!!!!
cfg.TEST.EVAL_PERIOD = 300 #number of iterations at which evaluation is run (Not relevant to validation loss calculation)
cfg.SOLVER.CHECKPOINT_PERIOD = 15000
cfg.DATALOADER.NUM_WORKERS = 2 #number of dataloading threads
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml") # Let training initialize from model zoo
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_101_FPN_3x.yaml") # Let training initialize from model zoo
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x.yaml") # Let training initialize from model zoo
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml") # Let training initialize from model zoo
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("Misc/scratch_mask_rcnn_R_50_FPN_3x_gn.yaml")
# cfg.MODEL.RESNETS.DEPTH = 101
# cfg.MODEL.RESNETS.RES2_OUT_CHANNELS = 64
cfg.MODEL.MASK_ON = False
cfg.MODEL.KEYPOINT_ON = True
cfg.MODEL.ROI_BOX_HEAD.SMOOTH_L1_BETA = 0.5 #Keypoint AP degrades (though box AP improves) when using plain L1 loss (i.e: value = 0.0)
cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 1500 #1000 proposals per-image is found to hurt box AP
cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 8 #needed since default number of keypoints is 17 in COCO dataset (for human pose estimation)
cfg.TEST.KEYPOINT_OKS_SIGMAS = [0.025, 0.025, 0.025, 0.025, 0.025, 0.025, 0.107, 0.107] #values show the importance of each keypoint location
#smaller=more precise - coco smallest and largest sigmas for human keypoints are used
#6th & 7th are assumed to be the usually hidden ones when having a vertical front face
# cfg.TEST.KEYPOINT_OKS_SIGMAS = [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.6]
cfg.SOLVER.IMS_PER_BATCH = 2 #1
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.MAX_ITER = 3000 #5000
cfg.SOLVER.GAMMA = 0.1 #lr decay factor (in multistep LR scheduler)
cfg.SOLVER.STEPS = [1000] #iteration milestones for reducing the lr (by gamma) #3000
cfg.SOLVER.WARMUP_FACTOR = 0 #start with a fraction of the learning rate for a number of iterations (warmup)
cfg.SOLVER.WARMUP_ITERS = 0 #warmup helps at initially avoiding learning irrelevant features
# cfg.SOLVER.NESTEROV = True
# cfg.SOLVER.WEIGHT_DECAY = 0
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256 #128 #number of ROIs to sample for training Fast RCNN head. sufficient for mini dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (cuboid)
cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS = True #False -> images without annotation are Not removed during training
cfg.MODEL.PIXEL_MEAN = [124.388, 121.619, 110.081] #BGR
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
my_trainer = MyTrainer(cfg)
# trainer = DefaultTrainer(cfg)
val_loss = ValidationLoss(cfg) #runs every 20 iterations by default (rate related to tensorboard writing)
my_trainer.register_hooks([val_loss])
# swap the order of PeriodicWriter and ValidationLoss
my_trainer._hooks = my_trainer._hooks[:-2] + my_trainer._hooks[-2:][::-1]
if train:
my_trainer.resume_or_load(resume=False)
my_trainer.train()
else:
my_trainer.resume_or_load(resume=True)
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
with open("./output/cfg_dump.txt", 'w') as file:
file.write(cfg.dump())
# print(cfg.dump())
return my_trainer, cfg
def load_image(image_path):
cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
object_classes = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes
cnn = EncoderCNN().to(device)
image_for_cnn = Image.open(image_path).convert("RGB")
image_for_cnn = transform(image_for_cnn)
image_features = cnn(image_for_cnn.unsqueeze(0).to(device))
image_for_predictor = cv2.imread(image_path)
object_detectection = predictor(image_for_predictor)
object_label_class_idxs = object_detectection["instances"].pred_classes
object_labels = list(
set([object_classes[item] for item in object_label_class_idxs]))
split_object_labels = []
for label in object_labels:
# a label can be multiple words (e.g. sports ball). We want to have these as individal words (e.g. ["sports", "ball"])
split_label = label.split()
split_object_labels.extend(split_label)
split_object_labels = list(set(split_object_labels))
return image_features, " ".join(split_object_labels)
def get_cfg(
outputn,
model="COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml",
iterations=1000,
):
dataset = CustomConfig.dataset
cfg = get_default()
cfg.merge_from_file(model_zoo.get_config_file(model))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model)
cfg.DATASETS.TRAIN = (dataset + "_train", )
cfg.DATASETS.TEST = (dataset + "_val", )
cfg.TEST.EVAL_PERIOD = 100
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = iterations
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = CustomConfig.numClasses
if isinstance(outputn, str):
cfg.OUTPUT_DIR = outputn
else:
cfg.OUTPUT_DIR = (f"./outputs/{CustomConfig.category}{outputn}")
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
print(f"Output Dir: {cfg.OUTPUT_DIR}")
return cfg
def get_masked_image(image):
cfg = get_cfg()
cfg.MODEL.DEVICE = 'cpu'
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(image)
predicted_classes = outputs["instances"].pred_classes.numpy()
if not (0 in predicted_classes):
return False
mask = outputs["instances"].pred_masks.numpy()[0].astype('uint8') * 255
#m3chan = cv2.merge((mask,mask,mask))
#masked = cv2.bitwise_and(image,m4chan)
#Transparency
b_channel, g_channel, r_channel = cv2.split(image)
#mask,_,_ = cv2.split(cropped_mask)
#alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255 #creating a dummy alpha channel image.
img_BGRA = cv2.merge((b_channel, g_channel, r_channel, mask))
#cv2_imshow(img_BGRA)
return img_BGRA
def new_detectron2_predictor(model_weights_file, NUM_CLASSES=2):
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
import torch, torchvision
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ('synth_train', )
cfg.DATASETS.TEST = ('synth_val', )
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.02
cfg.SOLVER.MAX_ITER = 300
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = NUM_CLASSES
cfg.MODEL.WEIGHTS = model_weights_file
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = 0.5
cfg.freeze()
return DefaultPredictor(cfg)
def add_imaterialist_config(cfg: CN):
"""
Add config for imaterialist2 head
"""
_C = cfg
_C.merge_from_file(model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
_C.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
##### Input #####
# Set a smaller image size than default to avoid memory problems
# Size of the smallest side of the image during training
# _C.INPUT.MIN_SIZE_TRAIN = (400,)
# # Maximum size of the side of the image during training
# _C.INPUT.MAX_SIZE_TRAIN = 600
# # Size of the smallest side of the image during testing. Set to zero to disable resize in testing.
# _C.INPUT.MIN_SIZE_TEST = 400
# # Maximum size of the side of the image during testing
# _C.INPUT.MAX_SIZE_TEST = 600
_C.SOLVER.IMS_PER_BATCH = 2
_C.SOLVER.BASE_LR = 0.0004
_C.SOLVER.MAX_ITER = 50000
_C.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 # default: 512
_C.MODEL.ROI_HEADS.NUM_CLASSES = 46 # 46 classes in iMaterialist
_C.MODEL.ROI_HEADS.NUM_ATTRIBUTES = 295
# this should ALWAYS be left at 1 because it will double or more memory usage if higher.
_C.DATALOADER.NUM_WORKERS = 1
def prepare_for_training(N_iter,
output_dir,
train_dataset_name,
N_classes,
start_training=False,
gpu_avail=True,
model_type="COCO-Detection/faster_rcnn_R_50_C4_1x.yaml"):
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(model_type))
cfg.OUTPUT_DIR = output_dir
cfg.DATASETS.TRAIN = (train_dataset_name,)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_type) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.GAMMA = 0.99 #lr decay
cfg.SOLVER.STEPS = list(range(1000, N_iter, 1000)) #(decay steps,)
cfg.SOLVER.WARMUP_ITERS = 500 #warmup steps
cfg.SOLVER.MAX_ITER = N_iter # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = N_classes # 4 classes
if not gpu_avail:
cfg.MODEL.DEVICE = 'cpu'
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
if start_training:
trainer.train()
return trainer, cfg
def get_local_cfg(val_fold=4, BATCH_SIZE=128):
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Detection/retinanet_R_101_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("lung_train", )
val_name = f"lung_val_{val_fold}"
cfg.best_model = f'best_{val_name}'
cfg.DATASETS.TEST = (val_name, )
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Detection/retinanet_R_101_FPN_3x.yaml"
) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon)
cfg.TEST.EVAL_PERIOD = 200
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = BATCH_SIZE # faster, and good enough for this toy dataset (default: 512)
cfg.SOLVER.MAX_ITER = 50000
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
return cfg
def __initEngine(self, score_threshold, model_zoo_config_path):
#initialize configuration
self.cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
self.cfg.merge_from_file(
model_zoo.get_config_file((model_zoo_config_path)))
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_threshold # score threshold to consider object of a class
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
model_zoo_config_path)
self.predictor = DefaultPredictor(self.cfg)
self.things = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0])
self.thing_classes = self.things.thing_classes
self.thing_colors = self.things.thing_colors
self.DEVICE = self.predictor.model.device.type + str(
self.predictor.model.device.index)
if self.selObjNames is None:
self.selObjNames = self.thing_classes
self.selObjIndices = list(range(len(self.selObjNames)))
else:
for n in self.selObjNames:
assert n in self.thing_classes, f"Error finding object class name: {n}"
self.selObjIndices.append(self.selObjNames.index(n))
def humanReplacement(filename):
cfg: CfgNode = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml")
predictor: DefaultPredictor = DefaultPredictor(cfg)
# Directory for taking and putting the image
directory = './static/client/img'
img: np.ndarray = cv2.imread('./static/client/origin_img.jpg')
changedImg: np.ndarray = cv2.imread(os.path.join(directory, filename))
x, y = img.shape[0:2]
changedImg = cv2.resize(changedImg, (y, x))
output: Instances = predictor(img)["instances"]
add_anime_characters(output, changedImg, filename)
v = Visualizer(img[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.0)
result: VisImage = v.draw_instance_predictions(output.to("cpu"))
result_image: np.ndarray = result.get_image()[:, :, ::-1]
cv2.imwrite('./static/client/mask_img.jpg', result_image)
def main(args):
# first regiest dataset I will use
register_coco_instances("my_dataset_train", {}, "training.json", "../data/ade20k/full_data/images/training/")
register_coco_instances("my_dataset_val", {}, "validation.json", "../data/ade20k/full_data/images/validation/")
# this is just a default cfg files
cfg = get_cfg()
# accordinig to different yaml file, it will change cfg files accordiningly
cfg.merge_from_file("../detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")
# This is some task specific changes I made for training ade20k dataset
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = () # no metrics implemented for this dataset
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 150 # 150 classes
cfg.SOLVER.IMS_PER_BATCH = 16 # this is default one
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
# I highly suggest read source code of DefaultTrainer again, if you forget why you did this.
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def get_modelzoo_config():
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(CFG_FILE))
#cfg.MODEL.WEIGHTS = r'/home/users/zzweng/unsupervised_segmentation/detectron2/checkpoints/model_final_5e3439.pkl'
#cfg.MODEL.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG = TRUE
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(CFG_FILE)
return cfg
def task_b_MOTS_training(model_name, model_file):
#model_name = model_name + '_inference'
print('Running task B for model', model_name)
SAVE_PATH = os.path.join('./results_week_5_task_b', model_name)
os.makedirs(SAVE_PATH, exist_ok=True)
# Load model and configuration
print('Loading Model')
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(model_file))
cfg.DATASETS.TRAIN = ('MOTS_train',)
cfg.DATASETS.TEST = ('KITTIMOTS_val',)
cfg.DATALOADER.NUM_WORKERS = 0
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.OUTPUT_DIR = SAVE_PATH
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_file)
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 200
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.TEST.SCORE_THRESH = 0.5
# Training
print('Training')
trainer = DefaultTrainer(cfg)
val_loss = ValidationLoss(cfg)
trainer.register_hooks([val_loss])
trainer._hooks = trainer._hooks[:-2] + trainer._hooks[-2:][::-1]
trainer.resume_or_load(resume=False)
trainer.train()
# Evaluation
print('Evaluating')
evaluator = COCOEvaluator('KITTIMOTS_val', cfg, False, output_dir=SAVE_PATH)
trainer.model.load_state_dict(val_loss.weights)
trainer.test(cfg, trainer.model, evaluators=[evaluator])
print('Plotting losses')
draw_loss(cfg, cfg.SOLVER.MAX_ITER, model_name, SAVE_PATH)
# Qualitative results: visualize some results
print('Getting qualitative results')
predictor = DefaultPredictor(cfg)
predictor.model.load_state_dict(trainer.model.state_dict())
inputs = kitti_val()
#inputs = inputs[:20] + inputs[-20:]
inputs = inputs[220:233] + inputs[1995:2100]
for i, input in enumerate(inputs):
file_name = input['file_name']
print('Prediction on image ' + file_name)
img = cv2.imread(file_name)
outputs = predictor(img)
v = Visualizer(
img[:, :, ::-1],
metadata=MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=0.8,
instance_mode=ColorMode.IMAGE)
v = v.draw_instance_predictions(outputs['instances'].to('cpu'))
cv2.imwrite(os.path.join(SAVE_PATH, 'Inference_' + model_name + '_inf_' + str(i) + '.png'), v.get_image()[:, :, ::-1])
def setup_detectron2_predictors(silhouettes_from='densepose'):
# Keypoint-RCNN
kprcnn_config_file = "COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"
kprcnn_cfg = get_cfg()
kprcnn_cfg.merge_from_file(model_zoo.get_config_file(kprcnn_config_file))
kprcnn_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model
kprcnn_cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(kprcnn_config_file)
kprcnn_cfg.freeze()
joints2D_predictor = DefaultPredictor(kprcnn_cfg)
if silhouettes_from == 'pointrend':
# PointRend-RCNN-R50-FPN
pointrend_config_file = "PointRend/configs/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco.yaml"
pointrend_cfg = get_cfg()
add_pointrend_config(pointrend_cfg)
pointrend_cfg.merge_from_file(pointrend_config_file)
pointrend_cfg.MODEL.WEIGHTS = "checkpoints/pointrend_rcnn_R_50_fpn.pkl"
pointrend_cfg.freeze()
silhouette_predictor = DefaultPredictor(pointrend_cfg)
elif silhouettes_from == 'densepose':
# DensePose-RCNN-R101-FPN
densepose_config_file = "DensePose/configs/densepose_rcnn_R_101_FPN_s1x.yaml"
densepose_cfg = get_cfg()
add_densepose_config(densepose_cfg)
densepose_cfg.merge_from_file(densepose_config_file)
densepose_cfg.MODEL.WEIGHTS = "checkpoints/densepose_rcnn_R_101_fpn_s1x.pkl"
densepose_cfg.freeze()
silhouette_predictor = DefaultPredictor(densepose_cfg)
return joints2D_predictor, silhouette_predictor
def make_cfg():
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
# cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
cfg.DATASETS.TRAIN = (_name, )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
) # Let training initialize from model zoo
# cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_0099999.pth")
# cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.SOLVER.IMS_PER_BATCH = 4
# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 150000 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (plane)
# user parameters
cfg.OUTPUT_DIR = "./output/holiCity_output"
cfg.INPUT.MASK_FORMAT = "bitmask" # polygon
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
return cfg
