This repository contains a Python implementation of Weighted Masks Fusion (WMF), which is a method for ensembling predicted masks from semantic or instance segmentation models. The WMF algorithm is inspirated by Weighted Boxes Fusion (WBF) introduced by ZFTurbo. It utilizes the principle of weighted fusion to combine masks from different segmentation models based on their confidence scores and other configurable parameters (e.g., conf_type, soft_weight, iou_thr, model_weights, etc.). The fused masks, scores, and boxes are used to generate an ensemble output with improved segmentation performance compared to individual models. The original code and concept for Weighted Masks Fusion were contributed by Kaggler Odede, whose solution achieved the 8th position in the "Sartorius - Cell Instance Segmentation" Kaggle competition.
Make sure you have the following prerequisites installed before using this library:
- Python 3
- Numba
- NumPy
You can install Weighted Mask Fusion by cloning this repository:
git clone https://github.com/chrise96/Weighted-Masks-Fusion.git
The following code example demonstrates how to use the weighted_masks_fusion
function to perform ensemble and fusion of masks from multiple YOLOv8 instance segmentation models.
inmodels = []
# Loop through all YOLOv8 models and perform predictions
for i, model in enumerate(yolo_models):
pred_boxes, pred_masks, scores, pred_classes = yolov8_predict(model, img)
inmodels += [i] * len(scores)
# Combine predictions from multiple models
if i == 0:
all_pred_boxes = pred_boxes
all_pred_masks = pred_masks
all_scores = scores
else:
all_pred_boxes = np.vstack([all_pred_boxes, pred_boxes])
all_pred_masks = np.vstack([all_pred_masks, pred_masks])
all_scores = np.hstack([all_scores, scores])
# Perform Weighted Masks Fusion on the combined predictions
pred_masks, scores, pred_boxes = weighted_masks_fusion(all_pred_masks, all_pred_boxes, all_scores, inmodels,
num_models=len(yolo_models), conf_type='model_weight', model_weights=[1,1,1]
)