Skip to content

Commit

Permalink
Merge branch 'main' into MMdet-Language-Vision-models-support
Browse files Browse the repository at this point in the history
  • Loading branch information
fcakyon authored Nov 22, 2024
2 parents d67bcb4 + de40562 commit 4163e0f
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 @@ -216,15 +216,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 4163e0f

Please sign in to comment.