From eba646d6816dbf3dde1dafa8e73b8b56b302d1f4 Mon Sep 17 00:00:00 2001
From: Shannon <22512825+lolipopshock@users.noreply.github.com>
Date: Tue, 28 Apr 2020 10:58:49 -0400
Subject: [PATCH] Improve class number calculation for COCO

---
 datasets/coco.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/datasets/coco.py b/datasets/coco.py
index c006f44..a25fa14 100644
--- a/datasets/coco.py
+++ b/datasets/coco.py
@@ -120,7 +120,7 @@ def image_aspect_ratio(self, image_index):
         return float(image['width']) / float(image['height'])
 
     def num_classes(self):
-        return 80
+        return len(self.coco.getCatIds())
 
 
 if __name__ == '__main__':