def __init__(self, cfg, is_train=True):
if cfg.INPUT.CROP.ENABLED and is_train:
self.crop_gen = [
T.ResizeShortestEdge([400, 500, 600], sample_style="choice"),
T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE),
]
else:
self.crop_gen = None
assert not cfg.MODEL.MASK_ON, "Mask is not supported"
self.tfm_gens = build_transform_gen(cfg, is_train)
logging.getLogger("detectron2").info(
"Full TransformGens used in pre-process: {}, crop: {}".format(
str(self.tfm_gens), str(self.crop_gen)))
self.img_format = cfg.INPUT.FORMAT
self.zip_read = cfg.INPUT.ZIP_READ
self.is_train = is_train
self.max_frame_dist = cfg.MODEL.DETR.MAX_FRAME_DIST
self.ignore_same = cfg.MODEL.DETR.IGNORE_SAME
# video2img --> 爆内存问题可能是出在这边，每次得释放掉 / 考虑放到builtin试试.
if self.is_train:
self.dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.
MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
if cfg.MODEL.LOAD_PROPOSALS else None,
)
assert len(cfg.DATASETS.TRAIN) == 1, logging.\
info("Only support ONE dataset each time, however, "
"there are {} datasets now.".format(len(cfg.DATASETS.TRAIN)))
self.video_to_images = MetadataCatalog.get(
cfg.DATASETS.TRAIN[0]).video_to_images
self.image_to_index = MetadataCatalog.get(
cfg.DATASETS.TRAIN[0]).image_to_index
else:
self.dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TEST,
filter_empty=False,
min_keypoints=0,
proposal_files=None)
assert len(cfg.DATASETS.TEST) == 1, logging. \
info("Only support ONE dataset each time, however, "
"there are {} datasets now.".format(len(cfg.DATASETS.TEST)))
self.video_to_images = MetadataCatalog.get(
cfg.DATASETS.TEST[0]).video_to_images
self.image_to_index = MetadataCatalog.get(
cfg.DATASETS.TEST[0]).image_to_index
def get_coco_test(batch_size, opts, cfg_file):
cfg = setup(opts, cfg_file)
dataset_name = 'coco_2017_val_panoptic_separated'
dataset_dicts = get_detection_dataset_dicts(
[dataset_name],
filter_empty=False,
proposal_files=None,
)
dataset = DatasetFromList(dataset_dicts)
mapper = DatasetMapper(cfg, False)
dataset = MapDataset(dataset, mapper)
sampler = samplers.InferenceSampler(len(dataset))
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler,
batch_size,
drop_last=False)
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=4,
batch_sampler=batch_sampler,
collate_fn=trivial_batch_collator,
)
return data_loader
def build_detection_test_loader(cfg, dataset_name, batch_size, mapper=None):
dataset_dicts = get_detection_dataset_dicts(
[dataset_name],
filter_empty=False,
proposal_files=[
cfg.DATASETS.PROPOSAL_FILES_TEST[list(
cfg.DATASETS.TEST).index(dataset_name)]
] if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset = DatasetFromList(dataset_dicts)
if mapper is None:
mapper = DatasetMapper(cfg, False)
dataset = MapDataset(dataset, mapper)
sampler = samplers.InferenceSampler(len(dataset))
# Always use 1 image per worker during inference since this is the
# standard when reporting inference time in papers.
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler,
batch_size,
drop_last=False)
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=cfg.DATALOADER.NUM_WORKERS,
batch_sampler=batch_sampler,
collate_fn=_trivial_batch_collator,
)
return data_loader
def get_coco_train(batch_size, opts, cfg_file):
cfg = setup(opts, cfg_file)
dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset = DatasetFromList(dataset_dicts, copy=False)
mapper = DatasetMapper(cfg, True)
dataset = MapDataset(dataset, mapper)
sampler = samplers.TrainingSampler(len(dataset))
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler,
batch_size,
drop_last=True)
# drop_last so the batch always have the same size
train_loader = torch.utils.data.DataLoader(
dataset,
num_workers=4,
batch_sampler=batch_sampler,
collate_fn=trivial_batch_collator,
worker_init_fn=worker_init_reset_seed,
)
return train_loader
def __init__(self,
dataset_dir=None,
model_pth_path=None,
score_threshold=0.15,
overlap_threshold=0.95,
model_config=None) -> None:
self.dataset_dir = dataset_dir
self.score_threshold = score_threshold
self.overlap_threshold = overlap_threshold
if model_config is None:
model_config = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
dataset_name = Path(self.dataset_dir).stem
self.subject_queue = queue.PriorityQueue(3)
self.left_object_queue = queue.PriorityQueue(3)
self.right_object_queue = queue.PriorityQueue(3)
self.right_interact_queue = queue.PriorityQueue(3)
self.left_interact_queue = queue.PriorityQueue(3)
self.subject_instance_name = 'Mouse'
self.left_object_name = 'LeftTeaball'
self.right_object_name = 'RightTeaball'
self.left_interact_name = 'LeftInteract'
self.right_interact_name = 'RightInteract'
self.tracking_results = []
try:
register_coco_instances(
f"{dataset_name}_train", {},
f"{self.dataset_dir}/train/annotations.json",
f"{self.dataset_dir}/train/")
register_coco_instances(
f"{dataset_name}_valid", {},
f"{self.dataset_dir}/valid/annotations.json",
f"{self.dataset_dir}/valid/")
except AssertionError as e:
print(e)
dataset_dicts = get_detection_dataset_dicts([f"{dataset_name}_train"])
_dataset_metadata = MetadataCatalog.get(f"{dataset_name}_train")
_dataset_metadata.thing_colors = [
cc['color'] for cc in builtin_meta.COCO_CATEGORIES
]
num_classes = len(_dataset_metadata.thing_classes)
self.class_names = _dataset_metadata.thing_classes
self.cfg = get_cfg()
# load model config and pretrained model
self.cfg.merge_from_file(model_zoo.get_config_file(model_config))
self.cfg.MODEL.WEIGHTS = model_pth_path
self.cfg.DATASETS.TRAIN = (f"{dataset_name}_train", )
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.score_threshold
self.cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes
self.cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = self.overlap_threshold
# NMS threshold used on RPN proposals
self.cfg.MODEL.RPN.NMS_THRESH = self.overlap_threshold
self.predictor = DefaultPredictor(self.cfg)
def build_gdrn_test_loader(cfg, dataset_name, train_objs=None):
"""Similar to `build_detection_train_loader`. But this function uses the
given `dataset_name` argument (instead of the names in cfg), and uses batch
size 1.
Args:
cfg: a detectron2 CfgNode
dataset_name (str): a name of the dataset that's available in the DatasetCatalog
Returns:
DataLoader: a torch DataLoader, that loads the given detection
dataset, with test-time transformation and batching.
"""
dataset_dicts = get_detection_dataset_dicts(
[dataset_name],
filter_empty=False,
proposal_files=[cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]]
if cfg.MODEL.LOAD_PROPOSALS
else None,
)
# load test detection results
if cfg.MODEL.LOAD_DETS_TEST:
det_files = cfg.DATASETS.DET_FILES_TEST
assert len(cfg.DATASETS.TEST) == len(det_files)
load_detections_into_dataset(
dataset_name,
dataset_dicts,
det_file=det_files[cfg.DATASETS.TEST.index(dataset_name)],
top_k_per_obj=cfg.DATASETS.DET_TOPK_PER_OBJ,
score_thr=cfg.DATASETS.DET_THR,
train_objs=train_objs,
)
if cfg.DATALOADER.FILTER_EMPTY_DETS:
dataset_dicts = filter_empty_dets(dataset_dicts)
dataset = GDRN_DatasetFromList(cfg, split="test", lst=dataset_dicts, flatten=False)
sampler = InferenceSampler(len(dataset))
# Always use 1 image per worker during inference since this is the
# standard when reporting inference time in papers.
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, 1, drop_last=False)
num_workers = cfg.DATALOADER.NUM_WORKERS
# Horovod: limit # of CPU threads to be used per worker.
# if num_workers > 0:
# torch.set_num_threads(num_workers)
kwargs = {"num_workers": num_workers}
# When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent
# issues with Infiniband implementations that are not fork-safe
# https://github.com/horovod/horovod/blob/master/examples/pytorch/pytorch_imagenet_resnet50.py
# if (num_workers > 0 and hasattr(mp, '_supports_context') and
# mp._supports_context and 'forkserver' in mp.get_all_start_methods()):
# kwargs['multiprocessing_context'] = 'forkserver'
data_loader = torch.utils.data.DataLoader(
dataset, batch_sampler=batch_sampler, collate_fn=trivial_batch_collator, **kwargs
)
return data_loader
def build_gdrn_train_loader(cfg, dataset_names):
"""A data loader is created by the following steps:
1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
2. Coordinate a random shuffle order shared among all processes (all GPUs)
3. Each process spawn another few workers to process the dicts. Each worker will:
* Map each metadata dict into another format to be consumed by the model.
* Batch them by simply putting dicts into a list.
The batched ``list[mapped_dict]`` is what this dataloader will yield.
Args:
cfg (CfgNode): the config
Returns:
an infinite iterator of training data
"""
dataset_dicts = get_detection_dataset_dicts(
dataset_names,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset_dicts = filter_invalid_in_dataset_dicts(
dataset_dicts, visib_thr=cfg.DATALOADER.FILTER_VISIB_THR)
dataset = GDRN_DatasetFromList(cfg,
split="train",
lst=dataset_dicts,
copy=False)
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
logger = logging.getLogger(__name__)
logger.info("Using training sampler {}".format(sampler_name))
# TODO avoid if-else?
if sampler_name == "TrainingSampler":
sampler = TrainingSampler(len(dataset))
elif sampler_name == "RepeatFactorTrainingSampler":
repeat_factors = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency(
dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD)
sampler = RepeatFactorTrainingSampler(repeat_factors)
else:
raise ValueError("Unknown training sampler: {}".format(sampler_name))
return my_build_batch_data_loader(
dataset,
sampler,
cfg.SOLVER.IMS_PER_BATCH,
aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING,
num_workers=cfg.DATALOADER.NUM_WORKERS,
)
def build_simple_dataloader(dataset_name: list, batch_size):
dataset_dicts = get_detection_dataset_dicts(dataset_name)
dataset = DatasetFromList(dataset_dicts, copy=False)
cfg = get_cfg()
cfg["aug_kwargs"] = {}
dataset = MapDataset(dataset, AlbumentationsMapper(cfg, False))
# set the shuffle to False in debugging mode
sampler = TrainingSampler(len(dataset), shuffle=False, seed=42)
dataloader = build_batch_data_loader(dataset=dataset, sampler=sampler,
total_batch_size=batch_size)
return dataloader
def build_detection_train_loader_with_train_sampler(cfg, mapper, seed=42, shuffle=True):
dataset_dicts = get_detection_dataset_dicts(cfg.DATASETS.TRAIN)
dataset = DatasetFromList(dataset_dicts, copy=False)
dataset = MapDataset(dataset, mapper)
logger = logging.getLogger(__name__)
logger.info("Using training sampler TrainingSampler with shuffle=False")
sampler = TrainingSampler(len(dataset), shuffle=shuffle, seed=seed)
return build_batch_data_loader(
dataset,
sampler,
cfg.SOLVER.IMS_PER_BATCH,
aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING,
num_workers=cfg.DATALOADER.NUM_WORKERS,
)
def build_weighted_detection_train_loader(cfg: CfgNode, mapper=None):
dataset_repeat_factors = get_train_datasets_repeat_factors(cfg)
# OrderedDict to guarantee order of values() consistent with repeat factors
dataset_name_to_dicts = OrderedDict(
{
name: get_detection_dataset_dicts(
[name],
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON
else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
if cfg.MODEL.LOAD_PROPOSALS
else None,
)
for name in cfg.DATASETS.TRAIN
}
)
# Repeat factor for every sample in the dataset
repeat_factors = [
[dataset_repeat_factors[dsname]] * len(dataset_name_to_dicts[dsname])
for dsname in cfg.DATASETS.TRAIN
]
repeat_factors = list(itertools.chain.from_iterable(repeat_factors))
dataset_dicts = dataset_name_to_dicts.values()
dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
dataset = DatasetFromList(dataset_dicts, copy=False)
if mapper is None:
mapper = DatasetMapper(cfg, True)
dataset = MapDataset(dataset, mapper)
logger.info(
"Using WeightedTrainingSampler with repeat_factors={}".format(
cfg.DATASETS.TRAIN_REPEAT_FACTOR
)
)
sampler = RepeatFactorTrainingSampler(torch.tensor(repeat_factors))
return build_batch_data_loader(
dataset,
sampler,
cfg.SOLVER.IMS_PER_BATCH,
aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING,
num_workers=cfg.DATALOADER.NUM_WORKERS,
)
def build_detection_test_loader(cfg, dataset_name, batch_size, mapper=None):
"""
Similar to `build_detection_train_loader`.
But this function uses the given `dataset_name` argument (instead of the names in cfg),
and uses batch size 1.
Args:
cfg: a detectron2 CfgNode
dataset_name (str): a name of the dataset that's available in the DatasetCatalog
mapper (callable): a callable which takes a sample (dict) from dataset
and returns the format to be consumed by the model.
By default it will be `DatasetMapper(cfg, False)`.
Returns:
DataLoader: a torch DataLoader, that loads the given detection
dataset, with test-time transformation and batching.
"""
dataset_dicts = get_detection_dataset_dicts(
[dataset_name],
filter_empty=False,
proposal_files=[
cfg.DATASETS.PROPOSAL_FILES_TEST[list(
cfg.DATASETS.TEST).index(dataset_name)]
] if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset = DatasetFromList(dataset_dicts)
if mapper is None:
mapper = DatasetMapper(cfg, False)
dataset = MapDataset(dataset, mapper)
sampler = samplers.InferenceSampler(len(dataset))
# Always use 1 image per worker during inference since this is the
# standard when reporting inference time in papers.
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler,
batch_size,
drop_last=False)
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=cfg.DATALOADER.NUM_WORKERS,
batch_sampler=batch_sampler,
collate_fn=trivial_batch_collator,
)
return data_loader
def main(args):
cfg = setup(args)
predictor = DefaultPredictor(cfg)
if args.img != '':
show_prediction(args.img, predictor, args.scale)
elif args.dir != '':
files = []
for ext in ['/*.jpg', '/*.png']:
files.extend(glob.glob(args.dir + ext))
for file in files:
key = show_prediction(file, predictor, args.scale)
if key == ord('q'):
cv2.destroyAllWindows()
break
else:
dataset_dicts = get_detection_dataset_dicts(['bizcard_val'])
for d in random.sample(dataset_dicts, 300):
key = show_prediction(d['file_name'], predictor, args.scale)
if key == ord('q'):
cv2.destroyAllWindows()
break
def build_train_loader(cfg, mapper=None):
if mapper is None:
mapper = get_dataset_mapper(cfg.DATASETS.TRAIN[0])(cfg, True)
dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset = DatasetFromList(dataset_dicts, copy=False)
dataset = MapDataset(dataset, mapper)
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
logger = logging.getLogger(__name__)
logger.info("Using training sampler {}".format(sampler_name))
# TODO avoid if-else?
if sampler_name == "TrainingSampler":
sampler = TrainingSampler(len(dataset))
elif sampler_name == "RepeatFactorTrainingSampler":
repeat_factors = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency(
dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD)
sampler = RepeatFactorTrainingSampler(repeat_factors)
elif sampler_name == "RatioFactorTrainingSampler":
repeat_factors = repeat_factors_from_ratios(dataset_dicts)
sampler = RepeatFactorTrainingSampler(repeat_factors)
else:
raise ValueError("Unknown training sampler: {}".format(sampler_name))
return build_batch_data_loader(
dataset,
sampler,
cfg.SOLVER.IMS_PER_BATCH,
aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING,
num_workers=cfg.DATALOADER.NUM_WORKERS,
)
#
# In[80]:
# In[81]:
register_coco_instances(f"{DATASET_NAME}_train", {},
f"{DATASET_DIR}/train/annotations.json",
f"{DATASET_DIR}/train/")
register_coco_instances(f"{DATASET_NAME}_valid", {},
f"{DATASET_DIR}/valid/annotations.json",
f"{DATASET_DIR}/valid/")
# In[25]:
dataset_dicts = get_detection_dataset_dicts([f"{DATASET_NAME}_train"])
# In[26]:
_dataset_metadata = MetadataCatalog.get(f"{DATASET_NAME}_train")
_dataset_metadata.thing_colors = [
cc['color'] for cc in builtin_meta.COCO_CATEGORIES
]
# In[27]:
_dataset_metadata
# In[28]:
NUM_CLASSES = len(_dataset_metadata.thing_classes)
def build_detection_train_loader(cfg, mapper=None):
"""
A data loader is created by the following steps:
1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
2. Coordinate a random shuffle order shared among all processes (all GPUs)
3. Each process spawn another few workers to process the dicts. Each worker will:
* Map each metadata dict into another format to be consumed by the model.
* Batch them by simply putting dicts into a list.
The batched ``list[mapped_dict]`` is what this dataloader will yield.
Args:
cfg (CfgNode): the config
mapper (callable): a callable which takes a sample (dict) from dataset and
returns the format to be consumed by the model.
By default it will be ``DatasetMapper(cfg, True)``.
Returns:
an infinite iterator of training data
"""
num_workers = get_world_size()
images_per_batch = cfg.SOLVER.IMS_PER_BATCH
assert (
images_per_batch % num_workers == 0
), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format(
images_per_batch, num_workers)
assert (
images_per_batch >= num_workers
), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format(
images_per_batch, num_workers)
images_per_worker = images_per_batch // num_workers
dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset = DatasetFromList(dataset_dicts, copy=False)
if mapper is None:
mapper = DatasetMapper(cfg, True)
dataset = MapDataset(dataset, mapper)
# d=dataset[0]
# print(d)
# print(d['image'].size())
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
logger = logging.getLogger(__name__)
logger.info("Using training sampler {}".format(sampler_name))
# TODO avoid if-else?
if sampler_name == "TrainingSampler":
sampler = TrainingSampler(len(dataset))
elif sampler_name == "RepeatFactorTrainingSampler":
repeat_factors = (
RepeatFactorTrainingSampler.repeat_factors_from_category_frequency(
dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD))
sampler = RepeatFactorTrainingSampler(repeat_factors)
elif sampler_name == "PairTrainingSampler":
sampler = PairTrainingSampler(cfg, dataset_dicts, images_per_worker)
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=cfg.DATALOADER.NUM_WORKERS,
batch_sampler=sampler,
collate_fn=trivial_batch_collator,
worker_init_fn=worker_init_reset_seed,
)
return PairDataLoader(cfg, data_loader)
else:
raise ValueError("Unknown training sampler: {}".format(sampler_name))
return build_batch_data_loader(
dataset,
sampler,
cfg.SOLVER.IMS_PER_BATCH,
aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING,
num_workers=cfg.DATALOADER.NUM_WORKERS,
)
cfg, args
) # if you don't like any of the default setup, write your own setup code
# if 'build_model(cfg)':
# meta_arch = cfg.MODEL.META_ARCHITECTURE
# model = META_ARCH_REGISTRY.get(meta_arch)(cfg)
# model.to(torch.device(cfg.MODEL.DEVICE))
if 'do-train':
...
if 'build_detection_train_loader':
# all dataset_dicts w.r.t cfg.DATASETS.TRAIN will be flattened.
dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON
else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
)
dataset = DatasetFromList(dataset_dicts, copy=False)
dataset = MapDataset(dataset, DatasetMapper(cfg, True))
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
logger = logging.getLogger(__name__)
logger.info("Using training sampler {}".format(sampler_name))
# TODO avoid if-else?
if sampler_name == "TrainingSampler":
sampler = TrainingSampler(len(dataset))
else:
raise ValueError("Unknown training sampler: {}".format(sampler_name))
data_loader = build_batch_data_loader(
def __init__(self,
dataset_dir=None,
out_put_dir=None,
score_threshold=0.15,
overlap_threshold=0.7,
max_iterations=3000,
batch_size=8,
model_config=None) -> None:
self.dataset_dir = dataset_dir
self.batch_size = batch_size
self.logger = logging.getLogger(__name__)
if model_config is None:
model_config = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
if out_put_dir is None:
self.out_put_dir = str(
Path(self.dataset_dir).parent / 'Annolid_training_outputs')
else:
self.out_put_dir = out_put_dir
self.score_threshold = score_threshold
self.overlap_threshold = overlap_threshold
self.dataset_name = Path(self.dataset_dir).stem
try:
register_coco_instances(
f"{self.dataset_name}_train", {},
f"{self.dataset_dir}/train/annotations.json",
f"{self.dataset_dir}/train/")
register_coco_instances(
f"{self.dataset_name}_valid", {},
f"{self.dataset_dir}/valid/annotations.json",
f"{self.dataset_dir}/valid/")
except AssertionError as e:
self.logger.info(e)
dataset_dicts = get_detection_dataset_dicts(
[f"{self.dataset_name}_train"])
_dataset_metadata = MetadataCatalog.get(f"{self.dataset_name}_train")
_dataset_metadata.thing_colors = [
cc['color'] for cc in builtin_meta.COCO_CATEGORIES
]
num_classes = len(_dataset_metadata.thing_classes)
self.class_names = _dataset_metadata.thing_classes
self.cfg = get_cfg()
# load model config and pretrained model
self.cfg.merge_from_file(model_zoo.get_config_file(model_config))
self.cfg.DATASETS.TRAIN = (f"{self.dataset_name}_train", )
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.score_threshold
self.cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes
self.cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = self.overlap_threshold
# NMS threshold used on RPN proposals
self.cfg.MODEL.RPN.NMS_THRESH = self.overlap_threshold
self.cfg.DATASETS.TEST = ()
self.cfg.DATALOADER.NUM_WORKERS = 2 # @param
self.cfg.DATALOADER.SAMPLER_TRAIN = "RepeatFactorTrainingSampler"
self.cfg.DATALOADER.REPEAT_THRESHOLD = 0.3
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
model_config) # Let training initialize from model zoo
self.cfg.SOLVER.IMS_PER_BATCH = self.batch_size # @param
self.cfg.SOLVER.BASE_LR = 0.0025 # @param # pick a good LR
self.logger.info(f"Max iterations {max_iterations}")
self.logger.info(f"Batch size is: {batch_size}")
self.logger.info(f"Dataset dir is : {dataset_dir}")
self.logger.info(f"Model config file is : {model_config}")
# @param # 3000 iterations seems good enough for 100 frames dataset;
# you will need to train longer for a practical dataset
self.cfg.SOLVER.MAX_ITER = max_iterations
self.cfg.SOLVER.CHECKPOINT_PERIOD = 1000 # @param
# @param # faster, and good enough for this toy dataset (default: 512)
self.cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
# (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
self.cfg.OUTPUT_DIR = self.out_put_dir
os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True)
self.trainer = DefaultTrainer(self.cfg)
self.trainer.resume_or_load(resume=True)
VizUtil.viz_mask(mask), mask],
axis=1)
cv2.imshow('view', view.astype(np.uint8))
key = cv2.waitKey(0)
if key == ord('q'):
cv2.destroyAllWindows()
return key
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--coco', action='store_true')
args = parser.parse_args()
if args.coco:
for d in random.sample(get_detection_dataset_dicts(['bizcard_train']),
300):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(
img[:, :, ::-1],
metadata=MetadataCatalog.get('bizcard_train'),
scale=0.2)
vis = visualizer.draw_dataset_dict(d)
cv2.imshow('', vis.get_image()[:, :, ::-1])
key = cv2.waitKey(0)
if key == ord('q'):
cv2.destroyAllWindows()
break
else:
view_gt_dir(
Path('/data/training/business_card/input/source_images'),
segmentations_dilated_flat, bbox_dilated_flat = process_mask_and_bbox(multi_poly_dilated_flat)
segmentations_dilated_square, bbox_dilated_square = process_mask_and_bbox(multi_poly_dilated_square)
save_file_with_new_masks(data, ann, img_id, segmentations_scaled, bbox_scaled, 'coco_dataset_train_visualizer_scale.json')
save_file_with_new_masks(data, ann, img_id, segmentations_dilated_round, bbox_dilated_round, 'coco_dataset_train_visualizer_dilated_round.json')
save_file_with_new_masks(data, ann, img_id, segmentations_dilated_flat, bbox_dilated_flat, 'coco_dataset_train_visualizer_dilated_flat.json')
save_file_with_new_masks(data, ann, img_id, segmentations_dilated_square, bbox_dilated_square, 'coco_dataset_train_visualizer_dilated_square.json')
root = Tk()
root.title("Image Viewer")
register_coco_instances("UCSP Fondecyt Dataset train", {}, "coco_dataset_train.json", "blurred_images")
register_coco_instances("UCSP Fondecyt Dataset train 2", {}, "coco_dataset_train.json", "new_images")
dataset_dicts = get_detection_dataset_dicts(["UCSP Fondecyt Dataset train"])
dataset_dicts2 = get_detection_dataset_dicts(["UCSP Fondecyt Dataset train 2"])
dataset_metadata = MetadataCatalog.get("UCSP Fondecyt Dataset train")
print(dataset_dicts[img_id-1]["file_name"])
my_img0 = ImageTk.PhotoImage(Image.open(dataset_dicts2[img_id-1]["file_name"][:-4] + "_w.png"))
my_img1 = ImageTk.PhotoImage(Image.open(dataset_dicts[img_id-1]["file_name"]))
#img = cv2.imread(dataset_dicts[img_id-1]["file_name"])
#visualizer = Visualizer(img[:, :, ::-1], metadata=dataset_metadata)
#vis = visualizer.draw_dataset_dict(dataset_dicts[img_id-1])
#my_img2 = ImageTk.PhotoImage(Image.fromarray(vis.get_image()))
img_orig = Image.open(dataset_dicts[img_id-1]["file_name"])
draw = ImageDraw.Draw(img_orig)
for list_coords in dataset_dicts[img_id-1]["annotations"][0]["segmentation"]:
draw.polygon(list_coords)
my_img2 = ImageTk.PhotoImage(img_orig)
