fix(setup): cpp extension (#1254)

Megvii-BaseDetection · Apr 18, 2022 · 1dee0f7 · 1dee0f7
1 parent 6513f76
commit 1dee0f7
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -8,6 +8,7 @@
 TORCH_AVAILABLE = True
 try:
     import torch
+    from torch.utils import cpp_extension
 except ImportError:
     TORCH_AVAILABLE = False
     print("[WARNING] Unable to import torch, pre-compiling ops will be disabled.")
@@ -56,7 +57,7 @@ def get_ext_modules():
 def get_cmd_class():
     cmdclass = {}
     if TORCH_AVAILABLE:
-        cmdclass["build_ext"] = torch.utils.cpp_extension.BuildExtension
+        cmdclass["build_ext"] = cpp_extension.BuildExtension
     return cmdclass
 
 

diff --git a/yolox/data/datasets/coco.py b/yolox/data/datasets/coco.py
@@ -191,7 +191,7 @@ def load_image(self, index):
         img_file = os.path.join(self.data_dir, self.name, file_name)
 
         img = cv2.imread(img_file)
-        assert img is not None
+        assert img is not None, f"file named {img_file} not found"
 
         return img
 

diff --git a/yolox/data/datasets/voc.py b/yolox/data/datasets/voc.py
@@ -219,7 +219,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)
-        assert img is not None
+        assert img is not None, f"file named {self._imgpath % img_id} not found"
 
         return img