keep info in coco dataset, which is necessary for newer pycocotools versions; also allow num_channels to be set in voc dataset

Author: lucia-leonie

yolox/data/datasets/coco.py

@@ -19,7 +19,6 @@ def remove_useless_info(coco):
     """
     if isinstance(coco, COCO):
         dataset = coco.dataset
-        dataset.pop("info", None)
         dataset.pop("licenses", None)
         for img in dataset["images"]:
             img.pop("license", None)
@@ -147,7 +146,7 @@ def load_image(self, index):
 
         img_file = os.path.join(self.data_dir, self.name, file_name)
 
-        img = cv2.imread(img_file, cv2.IMREAD_COLOR_BGR if self.num_channels == 3 else cv2.IMREAD_GRAYSCALE)
+        img = cv2.imread(img_file, cv2.IMREAD_COLOR if self.num_channels == 3 else cv2.IMREAD_GRAYSCALE)
         assert img is not None, f"file named {img_file} not found"
 
         return img

yolox/data/datasets/voc.py

@@ -108,6 +108,7 @@ def __init__(
         dataset_name="VOC0712",
         cache=False,
         cache_type="ram",
+        num_channels=3
     ):
         self.root = data_dir
         self.image_set = image_sets
@@ -131,6 +132,7 @@ def __init__(
             ):
                 self.ids.append((rootpath, line.strip()))
         self.num_imgs = len(self.ids)
+        self.num_channels = num_channels
 
         self.annotations = self._load_coco_annotations()
 
@@ -184,7 +186,7 @@ def load_resized_img(self, index):
 
     def load_image(self, index):
         img_id = self.ids[index]
-        img = cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)
+        img = cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR if self.num_channels == 3 else cv2.IMREAD_GRAYSCALE)
         assert img is not None, f"file named {self._imgpath % img_id} not found"
 
         return img