File size: 8,202 Bytes
b8597df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
import json
import os
import numpy as np
import pandas as pd
import torch
from pycocotools.coco import COCO
from torchvision.ops.boxes import box_convert, box_iou
from tqdm import tqdm
class NpEncoder(json.JSONEncoder):
"""Custom JSON encoder for NumPy data types.
This encoder handles NumPy-specific types that are not serializable by
the default JSON library by converting them into standard Python types.
"""
def default(self, obj):
"""Converts NumPy objects to their native Python equivalents.
Args:
obj (any): The object to encode.
Returns:
any: The JSON-serializable representation of the object.
"""
if isinstance(obj, np.integer):
return int(obj)
elif isinstance(obj, np.floating):
return float(obj)
elif isinstance(obj, np.ndarray):
return obj.tolist()
else:
return super(NpEncoder, self).default(obj)
class Ensembler:
"""A class to ensemble predictions from multiple object detection models.
This class loads ground truth data and predictions from several models,
performs non-maximum suppression (NMS) to merge overlapping detections,
and saves the final ensembled results in COCO format.
"""
def __init__(
self, output_dir, dataset_name, grplist, iou_thresh, coco_gt_path=None, coco_instances_results_fname=None
):
"""Initializes the Ensembler.
Args:
output_dir (str): The base directory where model outputs and
ensembled results are stored.
dataset_name (str): The name of the dataset being evaluated.
grplist (list[str]): A list of subdirectory names, where each
subdirectory contains the prediction file from one model.
iou_thresh (float): The IoU threshold for considering two bounding
boxes as overlapping during NMS.
coco_gt_path (str, optional): The full path to the ground truth
COCO JSON file. If None, it's assumed to be in `output_dir`.
Defaults to None.
coco_instances_results_fname (str, optional): The filename for the
COCO prediction files within each model's subdirectory.
Defaults to "coco_instances_results.json".
"""
self.output_dir = output_dir
self.dataset_name = dataset_name
self.grplist = grplist
self.iou_thresh = iou_thresh
self.n_detectors = len(grplist)
if coco_gt_path is None:
fname_gt = os.path.join(output_dir, dataset_name + "_coco_format.json")
else:
fname_gt = coco_gt_path
if coco_instances_results_fname is None:
fname_dt = "coco_instances_results.json"
else:
fname_dt = coco_instances_results_fname
# load in ground truth (form image lists)
coco_gt = COCO(fname_gt)
# populate detector truths
dtlist = []
for grp in grplist:
fname = os.path.join(output_dir, grp, fname_dt)
dtlist.append(coco_gt.loadRes(fname))
print("Successfully loaded {} into memory. {} instance detected.\n".format(fname, len(dtlist[-1].anns)))
self.coco_gt = coco_gt
self.cats = [cat["id"] for cat in self.coco_gt.dataset["categories"]]
self.dtlist = dtlist
self.results = []
print(
"Working with {} models, {} categories, and {} images.".format(
self.n_detectors, len(self.cats), len(self.coco_gt.imgs.keys())
)
)
def mean_score_nms(self):
"""Performs non-maximum suppression by merging overlapping boxes.
This method iterates through all images and categories, merging sets of
overlapping bounding boxes from different detectors based on the IoU
threshold. For each merged set, it calculates a mean score and selects
the single box with the highest original score as the representative
detection for the ensembled output.
Returns:
Ensembler: The instance itself, with the `self.results` attribute
populated with the ensembled predictions.
"""
def nik_merge(lsts):
"""Niklas B. https://github.com/rikpg/IntersectionMerge/blob/master/core.py"""
sets = [set(lst) for lst in lsts if lst]
merged = 1
while merged:
merged = 0
results = []
while sets:
common, rest = sets[0], sets[1:]
sets = []
for x in rest:
if x.isdisjoint(common):
sets.append(x)
else:
merged = 1
common |= x
results.append(common)
sets = results
return sets
winning_list = []
print(
"Computing mean score non-max suppression ensembling for {} images.".format(len(self.coco_gt.imgs.keys()))
)
for img in tqdm(self.coco_gt.imgs.keys()):
# print(img)
dflist = [] # a dataframe of detections
obj_set = set() # a set of objects (frozensets)
for i, coco_dt in enumerate(self.dtlist): # for each detector append predictions to df
dflist.append(pd.DataFrame(coco_dt.imgToAnns[img]).assign(det=i))
df = pd.concat(dflist, ignore_index=True)
if not df.empty:
for cat in self.cats: # for each category
dfcat = df[df["category_id"] == cat]
ts = box_convert(
torch.tensor(dfcat["bbox"]), in_fmt="xywh", out_fmt="xyxy"
) # list of tensor boxes for cateogory
iou_bool = np.array((box_iou(ts, ts) > self.iou_thresh)) # compute IoU matrix and threshold
for i in range(len(dfcat)): # for each detection in that category
fset = frozenset(dfcat.index[iou_bool[i]])
obj_set.add(fset) # compute set of sets representing objects
# find overlapping sets
# for fs in obj_set: #for existing sets
# if fs&fset: #check for
# fsnew = fs.union(fset)
# obj_set.remove(fs)
# obj_set.add(fsnew)
obj_set = nik_merge(obj_set)
for s in obj_set: # for each detected objects, find winning box and assign score as mean of scores
dfset = dfcat.loc[list(s)]
mean_score = dfset["score"].sum() / max(
self.n_detectors, len(s)
) # allows for more detections than detectors
winning_box = dfset.iloc[dfset["score"].argmax()].to_dict()
winning_box["score"] = mean_score
winning_list.append(winning_box)
print("{} resulting instances from NMS".format(len(winning_list)))
self.results = winning_list
return self
def save_coco_instances(self, fname="coco_instances_results.json"):
"""Saves the ensembled prediction results to a JSON file.
The output file follows the COCO instance format and can be used for
further evaluation.
Args:
fname (str, optional): The filename for the output JSON file.
Defaults to "coco_instances_results.json".
"""
if self.results:
with open(os.path.join(self.output_dir, fname), "w") as f:
f.write(json.dumps(self.results, cls=NpEncoder))
f.flush()
if __name__ == "__main__":
# Example usage:
# This assumes an 'output' directory with subdirectories 'fold1', 'fold2', etc.,
# each containing a 'coco_instances_results.json' file.
ens = Ensembler("dev", ["fold1", "fold2", "fold3", "fold4", "fold5"], 0.2)
ens.mean_score_nms()
|