Skip to content

Commit 7d3d9f8

Browse files
committed
Merge COCOReader changes needed for SSD correctness (#229)
Merged from wojciechp/dali-with-changes master branch Originally from #165 Co-authored-by: wojciechp@nvidia.com Reviewed-by: @slayton58 Signed-off-by: Cliff Woolley <jwoolley@nvidia.com>
1 parent b081a5f commit 7d3d9f8

File tree

1 file changed

+19
-1
lines changed

1 file changed

+19
-1
lines changed

dali/pipeline/operators/reader/coco_reader_op.h

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,11 +104,29 @@ class COCOReader : public DataReader<CPUBackend> {
104104
auto annotations = find_in_json(j, "annotations");
105105
int annotation_size = (*annotations).size();
106106

107+
// Change categories IDs to be in range [1, 80]
108+
std::vector<int> deleted_categories{ 12, 26, 29, 30, 45, 66, 68, 69, 71, 83, 91 };
109+
std::map<int, int> new_category_ids;
110+
int current_id = 1;
111+
int vector_id = 0;
112+
for (int i = 1; i <= 90; i++) {
113+
if (i == deleted_categories[vector_id]) {
114+
vector_id++;
115+
} else {
116+
new_category_ids.insert(std::make_pair(i, current_id));
117+
current_id++;
118+
}
119+
}
120+
107121
for (auto& an : *annotations) {
108122
auto image_id = get_from_json<int>(an, "image_id");
109123
auto category_id = get_from_json<int>(an, "category_id");
110124
auto bbox = get_from_json<std::array<float, 4>>(an, "bbox");
111125

126+
if (bbox[2] < 0.1 || bbox[3] < 0.1) {
127+
continue;
128+
}
129+
112130
if (ltrb_) {
113131
bbox[2] += bbox[0];
114132
bbox[3] += bbox[1];
@@ -127,7 +145,7 @@ class COCOReader : public DataReader<CPUBackend> {
127145

128146
annotations_multimap_.insert(
129147
std::make_pair(image_id,
130-
Annotation(bbox[0], bbox[1], bbox[2], bbox[3], category_id)));
148+
Annotation(bbox[0], bbox[1], bbox[2], bbox[3], new_category_ids[category_id])));
131149
}
132150

133151
f.close();

0 commit comments

Comments
 (0)