Skip to content

Commit

Permalink
Update has_mask method for mmdet models (handle ConcatDataset) (#1092)
Browse files Browse the repository at this point in the history
Co-authored-by: fatih c. akyon <[email protected]>
  • Loading branch information
ccomkhj and fcakyon authored Nov 22, 2024
1 parent 4981c57 commit de40562
Showing 1 changed file with 24 additions and 8 deletions.
32 changes: 24 additions & 8 deletions sahi/models/mmdet.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,15 +188,31 @@ def num_categories(self):
@property
def has_mask(self):
"""
Returns if model output contains segmentation mask
Returns if model output contains segmentation mask.
Considers both single dataset and ConcatDataset scenarios.
"""
# has_mask = self.model.model.with_mask
train_pipeline = self.model.cfg["train_dataloader"]["dataset"]["pipeline"]
has_mask = any(
isinstance(item, dict) and any("mask" in key and value is True for key, value in item.items())
for item in train_pipeline
)
return has_mask

def check_pipeline_for_mask(pipeline):
return any(
isinstance(item, dict) and any("mask" in key and value is True for key, value in item.items())
for item in pipeline
)

# Access the dataset from the configuration
dataset_config = self.model.cfg["train_dataloader"]["dataset"]

if dataset_config["type"] == "ConcatDataset":
# If using ConcatDataset, check each dataset individually
datasets = dataset_config["datasets"]
for dataset in datasets:
if check_pipeline_for_mask(dataset["pipeline"]):
return True
else:
# Otherwise, assume a single dataset with its own pipeline
if check_pipeline_for_mask(dataset_config["pipeline"]):
return True

return False

@property
def category_names(self):
Expand Down

0 comments on commit de40562

Please sign in to comment.