def _add_category_id_to_contiguous_id_maps_to_metadata(dataset_names: Iterable[str]):
# merge categories for all datasets
merged_categories = {}
for dataset_name in dataset_names:
meta = MetadataCatalog.get(dataset_name)
for cat_id, cat_name in meta.categories.items():
if cat_id not in merged_categories:
merged_categories[cat_id] = (cat_name, dataset_name)
continue
cat_name_other, dataset_name_other = merged_categories[cat_id]
if cat_name_other != cat_name:
raise ValueError(
f"Incompatible categories for category ID {cat_id}: "
f'dataset {dataset_name} value "{cat_name}", '
f'dataset {dataset_name_other} value "{cat_name_other}"'
)
merged_cat_id_to_cont_id = {}
for i, cat_id in enumerate(sorted(merged_categories.keys())):
merged_cat_id_to_cont_id[cat_id] = i
# add category maps to metadata
for dataset_name in dataset_names:
meta = MetadataCatalog.get(dataset_name)
categories = meta.get("categories")
meta.thing_classes = [categories[cat_id] for cat_id in sorted(categories.keys())]
meta.thing_dataset_id_to_contiguous_id = {
cat_id: merged_cat_id_to_cont_id[cat_id] for cat_id in sorted(categories.keys())
}
meta.thing_contiguous_id_to_dataset_id = {
merged_cat_id_to_cont_id[cat_id]: cat_id for cat_id in sorted(categories.keys())
}
def register_dacon_instances(name, metadata, json_file, image_root):
"""
Register a dataset in COCO's json annotation format for
instance detection, instance segmentation and keypoint detection.
(i.e., Type 1 and 2 in http://cocodataset.org/#format-data.
`instances*.json` and `person_keypoints*.json` in the dataset).
This is an example of how to register a new dataset.
You can do something similar to this function, to register new datasets.
Args:
name (str): the name that identifies a dataset, e.g. "coco_2014_train".
metadata (dict): extra metadata associated with this dataset. You can
leave it as an empty dict.
json_file (str): path to the json instance annotation file.
image_root (str or path-like): directory which contains all the images.
"""
assert isinstance(name, str), name
assert isinstance(json_file, (str, os.PathLike)), json_file
assert isinstance(image_root, (str, os.PathLike)), image_root
# 1. register a function which returns dicts
DatasetCatalog.register(
name, lambda: load_dacon_json(json_file, image_root, name))
# 2. Optionally, add metadata about this dataset,
# since they might be useful in evaluation, visualization or logging
MetadataCatalog.get(name).set(json_file=json_file,
image_root=image_root,
evaluator_type="coco",
**metadata)
def register_hico_instances(name, metadata, json_file, image_root,
evaluator_type):
"""
Register a hico-det dataset in COCO's json annotation format for human-object
interaction detection (i.e., `instances_hico_*.json` in the dataset).
This is an example of how to register a new dataset.
You can do something similar to this function, to register new datasets.
Args:
name (str): the name that identifies a dataset, e.g. "hico-det".
metadata (dict): extra metadata associated with this dataset. You can
leave it as an empty dict.
json_file (str): path to the json instance annotation file.
image_root (str or path-like): directory which contains all the images.
"""
# 1. register a function which returns dicts
DatasetCatalog.register(
name, lambda: load_hico_json(json_file, image_root, name))
# 2. Optionally, add metadata about this dataset,
# since they might be useful in evaluation, visualization or logging
MetadataCatalog.get(name).set(json_file=json_file,
image_root=image_root,
evaluator_type=evaluator_type,
**metadata)
def prepare_predictor():
class_list = ['diaodeng', 'pishafa'] ### 需要自己添加
train_metadata = 'zc_data' ### 需要自己添加
# current_path = os.getcwd()
current_path = os.path.abspath(__file__)
father_path = os.path.abspath(
os.path.dirname(current_path) + os.path.sep + ".")
# print("father_path:",father_path)
modelFile = os.path.join(father_path, "models/model20200728.pth")
cfgFile = os.path.join(
father_path,
"configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
num_class = len(class_list)
MetadataCatalog.get(train_metadata).thing_classes = class_list
train_metadata = MetadataCatalog.get(train_metadata)
# create config
cfg = get_cfg()
# below path applies to current installation location of Detectron2
cfg.merge_from_file(cfgFile)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set threshold for this model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_class
cfg.MODEL.WEIGHTS = modelFile
cfg.MODEL.DEVICE = "cuda" # we use a CPU Detectron copy
classes = train_metadata.get("thing_classes", None)
predictor = DefaultPredictor(cfg)
print("Predictor has been initialized.")
return (predictor, classes)
def prepare_pridctor():
class_list = ['diaodeng', 'pishafa'] ### 需要自己添加
train_metadata = 'zc_data' ### 需要自己添加
modelFile = "/home/user/qunosen/2_project/0_dvep/1_detectron2/ImageDetectionAPI/d2_object_detection/models/model20200728.pth"
cfgFile = "/home/user/qunosen/2_project/0_dvep/1_detectron2/ImageDetectionAPI/d2_object_detection/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
num_class = len(class_list)
MetadataCatalog.get(train_metadata).thing_classes = class_list
train_metadata = MetadataCatalog.get(train_metadata)
# create config
cfg = get_cfg()
# below path applies to current installation location of Detectron2
cfg.merge_from_file(cfgFile)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set threshold for this model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_class
cfg.MODEL.WEIGHTS = modelFile
cfg.MODEL.DEVICE = "cpu" # we use a CPU Detectron copy
classes = train_metadata.get("thing_classes", None)
predictor = DefaultPredictor(cfg)
print("Predictor has been initialized.")
return (predictor, classes)
def prepare_predictor():
current_path = os.path.abspath(__file__)
father_path = os.path.abspath(
os.path.dirname(current_path) + os.path.sep + ".")
# modelFile = os.path.join(father_path,"models/model20200728.pth")
modelFile = os.path.join(
father_path, "models/model_6class20200918_final.pth") ### class-6
cfgFile = os.path.join(
father_path,
"configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
num_class = len(class_list)
MetadataCatalog.get(regist_train_name).thing_classes = class_list
train_metadata = MetadataCatalog.get(regist_train_name)
# create config
cfg = get_cfg()
# below path applies to current installation location of Detectron2
cfg.merge_from_file(cfgFile)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set threshold for this model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_class
cfg.MODEL.WEIGHTS = modelFile
# cfg.MODEL.DEVICE = "cuda" # we use a GPU
cfg.MODEL.DEVICE = "cpu" # we use a GPU
classes = train_metadata.get("thing_classes", None)
predictor = DefaultPredictor(cfg)
print("Predictor has been initialized.")
return (predictor, classes)
def build_inference_based_loader(
cfg: CfgNode,
dataset_cfg: CfgNode,
model: torch.nn.Module,
embedder: Optional[torch.nn.Module] = None,
) -> InferenceBasedLoader:
"""
Constructs data loader based on inference results of a model.
"""
dataset = build_bootstrap_dataset(dataset_cfg.DATASET, dataset_cfg.IMAGE_LOADER)
meta = MetadataCatalog.get(dataset_cfg.DATASET)
training_sampler = TrainingSampler(len(dataset))
data_loader = torch.utils.data.DataLoader(
dataset, # pyre-ignore[6]
batch_size=dataset_cfg.IMAGE_LOADER.BATCH_SIZE,
sampler=training_sampler,
num_workers=dataset_cfg.IMAGE_LOADER.NUM_WORKERS,
collate_fn=trivial_batch_collator,
worker_init_fn=worker_init_reset_seed,
)
return InferenceBasedLoader(
model,
data_loader=data_loader,
data_sampler=build_data_sampler(cfg, dataset_cfg.DATA_SAMPLER, embedder),
data_filter=build_data_filter(dataset_cfg.FILTER),
shuffle=True,
batch_size=dataset_cfg.INFERENCE.OUTPUT_BATCH_SIZE,
inference_batch_size=dataset_cfg.INFERENCE.INPUT_BATCH_SIZE,
category_to_class_mapping=meta.category_to_class_mapping,
)
def get_detection_dataset_dicts_support(dataset_name,
filter_empty=True,
min_keypoints=0,
proposal_file=None):
dataset_dict = DatasetCatalog.get(dataset_name)
dataset_dict_flattened = []
for id_class, annotations_class in dataset_dict.items():
dataset_dict_flattened.extend(annotations_class)
# pre-extracted proposal: need to think about this later
"""
if proposal_files is not None:
assert len(dataset_names) == len(proposal_files)
# load precomputed proposals from proposal files
dataset_dicts = [
load_proposals_into_dataset(dataset_i_dicts, proposal_file)
for dataset_i_dicts, proposal_file in zip(dataset_dicts, proposal_files)
]
"""
# data distribution
class_names = MetadataCatalog.get(dataset_name).thing_classes
check_metadata_consistency("thing_classes", [dataset_name])
print_instances_class_histogram(dataset_dict_flattened, class_names)
return dataset_dict
def get_hoi_dataset_dicts(dataset_names, filter_empty=True):
"""
Load and prepare dataset dicts for HOI detection.
Args:
dataset_names (list[str]): a list of dataset names
filter_empty (bool): whether to filter out images without instance annotations
min_keypoints (int): filter out images with fewer keypoints than
`min_keypoints`. Set to 0 to do nothing.
proposal_files (list[str]): if given, a list of object proposal files
that match each dataset in `dataset_names`.
"""
assert len(dataset_names)
dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in dataset_names]
for dataset_name, dicts in zip(dataset_names, dataset_dicts):
assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
has_instances = "annotations" in dataset_dicts[0]
if filter_empty and has_instances:
dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
if filter_empty and has_instances and "actions" in dataset_dicts[0]["annotations"][0]:
dataset_dicts = filter_images_without_any_hois(dataset_dicts)
if has_instances:
try:
class_names = MetadataCatalog.get(dataset_names[0]).thing_classes
check_metadata_consistency("thing_classes", dataset_names)
print_instances_class_histogram(dataset_dicts, class_names)
except AttributeError: # class names are not available for this dataset
pass
return dataset_dicts
def build_bootstrap_dataset(dataset_name: str,
cfg: CfgNode) -> Sequence[torch.Tensor]:
"""
Build dataset that provides data to bootstrap on
Args:
dataset_name (str): Name of the dataset, needs to have associated metadata
to load the data
cfg (CfgNode): bootstrapping config
Returns:
Sequence[Tensor] - dataset that provides image batches, Tensors of size
[N, C, H, W] of type float32
"""
logger = logging.getLogger(__name__)
_add_category_info_to_bootstrapping_metadata(dataset_name, cfg)
meta = MetadataCatalog.get(dataset_name)
factory = BootstrapDatasetFactoryCatalog.get(meta.dataset_type)
dataset = None
if factory is not None:
dataset = factory(meta, cfg)
if dataset is None:
logger.warning(
f"Failed to create dataset {dataset_name} of type {meta.dataset_type}"
)
return dataset
def load_dacon_test_json(json_file,
image_root,
dataset_name=None,
extra_annotation_keys=None):
timer = Timer()
json_file = PathManager.get_local_path(json_file)
with contextlib.redirect_stdout(io.StringIO()):
dacon_api = DaconAPI(json_file)
anns = dacon_api.features
if timer.seconds() > 1:
logger.info("Loading {} takes {:.2f} seconds.".format(
json_file, timer.seconds()))
meta = MetadataCatalog.get(dataset_name)
meta.thing_classes = dacon_api.thing_classes
logger.info("Loaded {} images in dacon format from {}".format(
len(anns), json_file))
dataset_dicts = []
for ann in anns:
record = {}
record["file_name"] = os.path.join(image_root, ann['image_id'])
dataset_dicts.append(record)
return dataset_dicts
def check_metadata_consistency(key, dataset_names):
"""
Check that the datasets have consistent metadata.
Args:
key (str): a metadata key
dataset_names (list[str]): a list of dataset names
Raises:
AttributeError: if the key does not exist in the metadata
ValueError: if the given datasets do not have the same metadata values defined by key
"""
if len(dataset_names) == 0:
return
logger = logging.getLogger(__name__)
print('-----', dataset_names)
entries_per_dataset = [
getattr(MetadataCatalog.get(d), key) for d in dataset_names
]
for idx, entry in enumerate(entries_per_dataset):
if entry != entries_per_dataset[0]:
logger.error("Metadata '{}' for dataset '{}' is '{}'".format(
key, dataset_names[idx], str(entry)))
logger.error("Metadata '{}' for dataset '{}' is '{}'".format(
key, dataset_names[0], str(entries_per_dataset[0])))
raise ValueError(
"Datasets have different metadata '{}'!".format(key))
def _map_category_id_to_contiguous_id(dataset_name: str,
dataset_dicts: Iterable[Instance]):
meta = MetadataCatalog.get(dataset_name)
for dataset_dict in dataset_dicts:
for ann in dataset_dict["annotations"]:
ann["category_id"] = meta.thing_dataset_id_to_contiguous_id[
ann["category_id"]]
def get_classification_dataset_dicts(dataset_names):
"""
Load and prepare dataset dicts for instance detection/segmentation and semantic segmentation.
Args:
dataset_names (list[str]): a list of dataset names
"""
assert len(dataset_names)
dataset_dicts = []
for dataset_name in dataset_names:
assert dataset_name in _GRUUL_DATASETS, f"Unknown train dataset {dataset_name}"
# dataset = ImageFolder(_GRUUL_DATASETS[cfg.DATASETS.TRAIN[0]])
dataset_dicts = [
DatasetCatalog.get(dataset_name) for dataset_name in dataset_names
]
for dataset_name, dicts in zip(dataset_names, dataset_dicts):
assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
try:
class_names = MetadataCatalog.get(dataset_names[0]).classes
print_instances_class_histogram(dataset_dicts, class_names)
except AttributeError: # class names are not available for this dataset
pass
return dataset_dicts
def _add_category_id_to_contiguous_id_maps_to_metadata(merged_categories: _MergedCategoriesT):
merged_categories_per_dataset = {}
for contiguous_cat_id, cat_id in enumerate(sorted(merged_categories.keys())):
for cat in merged_categories[cat_id]:
if cat.dataset_name not in merged_categories_per_dataset:
merged_categories_per_dataset[cat.dataset_name] = defaultdict(list)
merged_categories_per_dataset[cat.dataset_name][cat_id].append(
(
contiguous_cat_id,
cat,
)
)
logger = logging.getLogger(__name__)
for dataset_name, merged_categories in merged_categories_per_dataset.items():
meta = MetadataCatalog.get(dataset_name)
if not hasattr(meta, "thing_classes"):
meta.thing_classes = []
meta.thing_dataset_id_to_contiguous_id = {}
meta.thing_dataset_id_to_merged_id = {}
else:
meta.thing_classes.clear()
meta.thing_dataset_id_to_contiguous_id.clear()
meta.thing_dataset_id_to_merged_id.clear()
logger.info(f"Dataset {dataset_name}: category ID to contiguous ID mapping:")
for _cat_id, categories in sorted(merged_categories.items()):
added_to_thing_classes = False
for contiguous_cat_id, cat in categories:
if not added_to_thing_classes:
meta.thing_classes.append(cat.mapped_name)
added_to_thing_classes = True
meta.thing_dataset_id_to_contiguous_id[cat.id] = contiguous_cat_id
meta.thing_dataset_id_to_merged_id[cat.id] = cat.mapped_id
logger.info(f"{cat.id} ({cat.name}) -> {contiguous_cat_id}")
def build_bootstrap_dataset(dataset_name: str,
cfg: CfgNode) -> Sequence[torch.Tensor]:
"""
Build dataset that provides data to bootstrap on
Args:
dataset_name (str): Name of the dataset, needs to have associated metadata
to load the data
cfg (CfgNode): bootstrapping config
Returns:
Sequence[Tensor] - dataset that provides image batches, Tensors of size
[N, C, H, W] of type float32
"""
logger = logging.getLogger(__name__)
meta = MetadataCatalog.get(dataset_name)
if meta.dataset_type == DatasetType.VIDEO_LIST:
video_list_fpath = meta.video_list_fpath
video_base_path = meta.video_base_path
if cfg.TYPE == "video_keyframe":
frame_selector = build_frame_selector(cfg.SELECT)
transform = build_transform(cfg.TRANSFORM, data_type="image")
video_list = video_list_from_file(video_list_fpath,
video_base_path)
return VideoKeyframeDataset(video_list, frame_selector, transform)
logger.warning(
f"Could not create an image sampler of type {cfg.TYPE} for dataset {dataset_name}"
)
return None
def from_config(cls, cfg, input_shape):
# fmt: on
num_classes = cfg.MODEL.ROI_HEADS.NUM_ACTIONS
test_score_thresh = cfg.MODEL.ROI_HEADS.HOI_SCORE_THRESH_TEST
test_topk_per_image = cfg.TEST.INTERACTIONS_PER_IMAGE
action_cls_weights = cfg.MODEL.HOI_BOX_HEAD.ACTION_CLS_WEIGHTS
batch_size_per_image = cfg.MODEL.ROI_HEADS.HOI_BATCH_SIZE_PER_IMAGE
ims_per_batch = cfg.SOLVER.IMS_PER_BATCH
# fmt: off
if cfg.MODEL.ROI_HEADS.POSITION_FEAT == 0:
input_dup = 3
elif cfg.MODEL.ROI_HEADS.POSITION_FEAT == 4:
input_dup = 6
else:
input_dup = 5
additional_dim = 0
if cfg.MODEL.ROI_HEADS.ADD_WORD_EMB == 1:
additional_dim = 300 * 2
if cfg.MODEL.ROI_HEADS.ADD_WORD_EMB == 2:
additional_dim = 300 * 3
from .roi_heads import get_convert_matrix
verb_to_HO_matrix, obj_to_HO_matrix = get_convert_matrix(
obj_class_num=81)
verb_to_HO_matrix = torch.from_numpy(verb_to_HO_matrix)
# Positive weights are used to balance the instances at training.
# Get the prior distribution from the metadata.
pos_weights = torch.full((num_classes, ), 0.)
for dataset in cfg.DATASETS.TRAIN:
meta = MetadataCatalog.get(dataset)
priors = meta.get("action_priors", None)
if priors and num_classes != 600:
priors = torch.as_tensor(
priors) * ims_per_batch * batch_size_per_image
pos_weights_per_dataset = torch.clamp(
1. / priors,
min=action_cls_weights[0],
max=action_cls_weights[1],
)
pos_weights += pos_weights_per_dataset
elif priors and num_classes == 600:
pos_weights = torch.from_numpy(
np.sqrt(sum(inst_nums) / inst_nums))
pass
else:
pos_weights += torch.full((num_classes, ), 1.)
if len(cfg.DATASETS.TRAIN):
pos_weights /= len(cfg.DATASETS.TRAIN)
return {
"input_shape": input_shape,
"num_classes": num_classes,
"pos_weights": pos_weights,
"test_score_thresh": test_score_thresh,
"test_topk_per_image": test_topk_per_image,
"input_dup": input_dup,
"additional_dim": additional_dim
}
def _add_category_whitelists_to_metadata(cfg: CfgNode):
for dataset_name, whitelisted_cat_ids in cfg.DATASETS.WHITELISTED_CATEGORIES.items(
):
meta = MetadataCatalog.get(dataset_name)
meta.whitelisted_categories = whitelisted_cat_ids
logger = logging.getLogger(__name__)
logger.info("Whitelisted categories for dataset {}: {}".format(
dataset_name, meta.whitelisted_categories))
def _add_category_info_to_bootstrapping_metadata(dataset_name: str,
dataset_cfg: CfgNode) -> None:
meta = MetadataCatalog.get(dataset_name)
meta.category_to_class_mapping = get_category_to_class_mapping(dataset_cfg)
meta.categories = dataset_cfg.CATEGORIES
meta.max_count_per_category = dataset_cfg.MAX_COUNT_PER_CATEGORY
logger = logging.getLogger(__name__)
logger.info("Category to class mapping for dataset {}: {}".format(
dataset_name, meta.category_to_class_mapping))
def _add_category_maps_to_metadata(cfg: CfgNode) -> None:
for dataset_name, category_map in cfg.DATASETS.CATEGORY_MAPS.items():
category_map = {
int(cat_id_src): int(cat_id_dst) for cat_id_src, cat_id_dst in category_map.items()
}
meta = MetadataCatalog.get(dataset_name)
meta.category_map = category_map
logger = logging.getLogger(__name__)
logger.info("Category maps for dataset {}: {}".format(dataset_name, meta.category_map))
def setup_cfg(args):
MetadataCatalog.get("custom")
cfg = get_cfg()
# cuda context is initialized before creating dataloader, so we don't fork anymore
cfg.DATALOADER.NUM_WORKERS = 0
cfg = add_export_config(cfg)
# cfg.merge_from_file(args.config_file)
cfg.merge_from_file(
# "../detectron2-master/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml" # COCO
"../../detectron2-master/configs/PascalVOC-Detection/faster_rcnn_R_50_FPN.yaml" # VOC
)
cfg.merge_from_list(args.opts)
cfg.MODEL.WEIGHTS = "../output/model_final.pth"
cfg.DATASETS.TRAIN = ("custom",)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4 # num classes
cfg.freeze()
if cfg.MODEL.DEVICE != "cpu":
assert TORCH_VERSION >= (1, 5), "PyTorch>=1.5 required for GPU conversion!"
return cfg
def register_vcoco_instances(name, metadata, json_file, image_root,
evaluator_type):
"""
Register a vcoco dataset in COCO's json annotation format for human-object
interaction detection (i.e., `instances_hico_*.json` in the dataset).
Args:
see `register_hico_instances`
"""
# 1. register a function which returns dicts
DatasetCatalog.register(
name, lambda: load_vcoco_json(json_file, image_root, name))
# 2. Optionally, add metadata about this dataset,
# since they might be useful in evaluation, visualization or logging
MetadataCatalog.get(name).set(json_file=json_file,
image_root=image_root,
evaluator_type=evaluator_type,
**metadata)
def print_label(self, pred_instances):
from detectron2.data.catalog import MetadataCatalog
catalog = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0])
labels = catalog.get("thing_classes")
for idx, img in enumerate(pred_instances):
instances = img["instances"]
classes = [labels[i] for i in instances.pred_classes]
print(f"image contains {', '.join(classes)}")
def get_detection_dataset_dicts(dataset_names,
filter_empty=True,
min_keypoints=0,
proposal_files=None):
"""
Load and prepare dataset dicts for instance detection/segmentation and semantic segmentation.
Args:
dataset_names (list[str]): a list of dataset names
filter_empty (bool): whether to filter out images without instance annotations
min_keypoints (int): filter out images with fewer keypoints than
`min_keypoints`. Set to 0 to do nothing.
proposal_files (list[str]): if given, a list of object proposal files
that match each dataset in `dataset_names`.
"""
assert len(dataset_names)
print(dataset_names)
dataset_dicts = [
DatasetCatalog.get(dataset_name) for dataset_name in dataset_names
]
for dataset_name, dicts in zip(dataset_names, dataset_dicts):
assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
if proposal_files is not None:
assert len(dataset_names) == len(proposal_files)
# load precomputed proposals from proposal files
dataset_dicts = [
load_proposals_into_dataset(dataset_i_dicts, proposal_file)
for dataset_i_dicts, proposal_file in zip(dataset_dicts,
proposal_files)
]
dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
has_instances = "annotations" in dataset_dicts[0]
# Keep images without instance-level GT if the dataset has semantic labels.
if filter_empty and has_instances and "sem_seg_file_name" not in dataset_dicts[
0]:
dataset_dicts = filter_images_with_only_crowd_annotations(
dataset_dicts)
if min_keypoints > 0 and has_instances:
dataset_dicts = filter_images_with_few_keypoints(
dataset_dicts, min_keypoints)
if has_instances:
try:
class_names = MetadataCatalog.get(dataset_names[0]).thing_classes
check_metadata_consistency("thing_classes", dataset_names)
print_instances_class_histogram(dataset_dicts, class_names)
plot_instances_class_histogram(dataset_dicts, class_names,
dataset_name)
except AttributeError: # class names are not available for this dataset
pass
return dataset_dicts
def combine_detection_dataset_dicts(
dataset_names: Collection[str],
keep_instance_predicate: Optional[InstancePredicate] = None,
proposal_files: Optional[Collection[str]] = None,
) -> List[Instance]:
"""
Load and prepare dataset dicts for training / testing
Args:
dataset_names (Collection[str]): a list of dataset names
keep_instance_predicate (Callable: Dict[str, Any] -> bool): predicate
applied to instance dicts which defines whether to keep the instance
proposal_files (Collection[str]): if given, a list of object proposal files
that match each dataset in `dataset_names`.
"""
assert len(dataset_names)
if proposal_files is None:
proposal_files = [None] * len(dataset_names)
assert len(dataset_names) == len(proposal_files)
# load datasets and metadata
dataset_name_to_dicts = {}
for dataset_name in dataset_names:
dataset_name_to_dicts[dataset_name] = DatasetCatalog.get(dataset_name)
assert len(
dataset_name_to_dicts), f"Dataset '{dataset_name}' is empty!"
# merge categories, requires category metadata to be loaded
# cat_id -> [(orig_cat_id, cat_name, dataset_name)]
merged_categories = _merge_categories(dataset_names)
_warn_if_merged_different_categories(merged_categories)
# map to contiguous category IDs
_add_category_id_to_contiguous_id_maps_to_metadata(merged_categories)
# load annotations and dataset metadata
for dataset_name, proposal_file in zip(dataset_names, proposal_files):
dataset_dicts = dataset_name_to_dicts[dataset_name]
assert len(dataset_dicts), f"Dataset '{dataset_name}' is empty!"
if proposal_file is not None:
dataset_dicts = load_proposals_into_dataset(
dataset_dicts, proposal_file)
dataset_dicts = _maybe_filter_and_map_categories(
dataset_name, dataset_dicts)
print_instances_class_histogram(
dataset_dicts,
MetadataCatalog.get(dataset_name).thing_classes)
dataset_name_to_dicts[dataset_name] = dataset_dicts
if keep_instance_predicate is not None:
all_datasets_dicts_plain = [
d for d in itertools.chain.from_iterable(
dataset_name_to_dicts.values()) if keep_instance_predicate(d)
]
else:
all_datasets_dicts_plain = list(
itertools.chain.from_iterable(dataset_name_to_dicts.values()))
return all_datasets_dicts_plain
def Predictor(yaml_file, weight_file, confidence_thr, num_classes):
"""
Input:
yaml_file: str. backbone
weight_file : str. model weight
confidence_thr: Num. confidence threshold
num_classes: Num. classes number
Output:
return : predictor
"""
MetadataCatalog.get("custom")
cfg = get_cfg()
cfg.merge_from_file(yaml_file)
cfg.MODEL.WEIGHTS = weight_file
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_thr
cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes
cfg.DATASETS.TEST = ()
predictor = DefaultPredictor(cfg)
return predictor
def __init__(self, cfg):
self.cfg = cfg.clone() # cfg can be modified by model
self.model = build_model(self.cfg)
self.model.eval()
self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
checkpointer = DetectionCheckpointer(self.model)
checkpointer.load(cfg.MODEL.WEIGHTS)
self.input_format = cfg.INPUT.FORMAT
assert self.input_format in ["RGB", "BGR"], self.input_format
def __init__(self, cfg, instance_mode=ColorMode.IMAGE):
"""
Args:
cfg (CfgNode):
instance_mode (ColorMode):
"""
self.metadata = MetadataCatalog.get(
cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused")
self.cpu_device = torch.device("cpu")
self.instance_mode = instance_mode
self.predictor = BatchPredictor(cfg)
def get_evaluator(cfg, dataset_name, output_folder=None):
evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
if evaluator_type == "lvis":
return LVISEvaluatorMaxDets(
dataset_name, cfg, True, output_folder, max_dets=[-1]
)
elif evaluator_type == "coco":
return COCOEvaluatorMaxDets(
dataset_name, cfg, True, output_folder, max_dets=[-1]
)
else:
raise NotImplementedError
def __init__(self, cfg, is_train=True):
if cfg.INPUT.CROP.ENABLED and is_train:
self.crop_gen = T.RandomCrop(cfg.INPUT.CROP.TYPE,
cfg.INPUT.CROP.SIZE)
logging.getLogger(__name__).info("CropGen used in training: " +
str(self.crop_gen))
else:
self.crop_gen = None
self.tfm_gens = utils.build_transform_gen(cfg, is_train)
# fmt: off
self.img_format = cfg.INPUT.FORMAT
self.mask_on = cfg.MODEL.MASK_ON
self.mask_format = cfg.INPUT.MASK_FORMAT
self.keypoint_on = cfg.MODEL.KEYPOINT_ON
self.load_proposals = cfg.MODEL.LOAD_PROPOSALS
self.few_shot = cfg.INPUT.FS.FEW_SHOT
self.support_way = cfg.INPUT.FS.SUPPORT_WAY
self.support_shot = cfg.INPUT.FS.SUPPORT_SHOT
# fmt: on
if self.keypoint_on and is_train:
# Flip only makes sense in training
self.keypoint_hflip_indices = utils.create_keypoint_hflip_indices(
cfg.DATASETS.TRAIN)
else:
self.keypoint_hflip_indices = None
if self.load_proposals:
self.proposal_min_box_size = cfg.MODEL.PROPOSAL_GENERATOR.MIN_SIZE
self.proposal_topk = (cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TRAIN
if is_train else
cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TEST)
self.is_train = is_train
if self.is_train:
# support_df
self.support_on = True
if self.few_shot:
self.support_df = pd.read_pickle(
"./datasets/coco/10_shot_support_df.pkl")
else:
self.support_df = pd.read_pickle(
"./datasets/coco/train_support_df.pkl")
metadata = MetadataCatalog.get('coco_2017_train')
# unmap the category mapping ids for COCO
reverse_id_mapper = lambda dataset_id: metadata.thing_dataset_id_to_contiguous_id[
dataset_id] # noqa
self.support_df['category_id'] = self.support_df[
'category_id'].map(reverse_id_mapper)
