Using Roboflow Universe datasets for training detection, segmentation and classification

.github/workflows/ci.yml

@@ -49,6 +49,9 @@ jobs:
           python -m unittest
 
       - name: Test scripts
+        env:
+          ROBOFLOW_TOKEN: ${{ secrets.ROBOFLOW_API_KEY }}
+        shell: bash  # for Windows compatibility
         run: |
           pip install -e .
           # donwload coco formatted testing data
@@ -89,3 +92,8 @@ jobs:
           yolov5 segment train --img 128 --weights fcakyon/yolov5n-seg-v7.0 --epochs 1 --device cpu
           # yolov5 segment val --img 128 --weights yolov5n-seg.pt --device cpu
           yolov5 segment predict --img 128 --weights yolov5n-seg.pt --device cpu
+          # roboflow
+          yolov5 train --data https://universe.roboflow.com/jacob-solawetz/aerial-maritime/dataset/10 --weights yolov5/weights/yolov5n.pt --device cpu --roboflow_token ${{ env.ROBOFLOW_TOKEN }}
+          yolov5 classify train --data https://universe.roboflow.com/bagas-etcfr/flowerdata1/dataset/1 --img 128 --model yolov5n-cls.pt --epochs 1 --device cpu --roboflow_token ${{ env.ROBOFLOW_TOKEN }}
+          yolov5 segment train --data https://universe.roboflow.com/scaffolding/scaffolding-distance/dataset/2 --img 128 --weights yolov5n-seg.pt --epochs 1 --device cpu --roboflow_token ${{ env.ROBOFLOW_TOKEN }}
+

README.md

@@ -165,6 +165,14 @@ val_image_dir: "val_image_dir/"
 $ yolov5 train --data data.yaml --weights yolov5s.pt
 ```
 
+- Train your model using [Roboflow Universe](https://universe.roboflow.com/) datasets (roboflow>=0.2.27 required):
+
+```bash
+$ yolov5 train --data DATASET_UNIVERSE_URL --weights yolov5s.pt --roboflow_token YOUR_ROBOFLOW_TOKEN
+```
+
+Where `DATASET_UNIVERSE_URL` must be in `https://universe.roboflow.com/workspace_name/project_name/project_version` format.
+
 - Visualize your experiments via [Neptune.AI](https://neptune.ai/) (neptune-client>=0.10.10 required):
 
 ```bash

requirements.txt

@@ -45,7 +45,7 @@ seaborn>=0.11.0
 # mss  # screenshots
 # albumentations>=1.0.3
 # pycocotools>=2.0  # COCO mAP
-# roboflow
+# roboflow>=0.2.27
 # ultralytics  # HUB https://hub.ultralytics.com
 
 # CLI

tests/test_yolov5.py

@@ -4,120 +4,8 @@
 
 DEVICE = "cpu"
 
-class TestYolov5FromUltralytics(unittest.TestCase):
+class TestYolov5FromUltralytics(unittest.TestCase):       
     def test_load_model(self):
-        from yolov5 import YOLOv5
-
-        # init model
-        model_path = TestConstants.YOLOV5S_MODEL_PATH
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=False)
-        yolov5.load_model()
-
-        # check if loaded
-        self.assertNotEqual(yolov5.model, None)
-
-    def test_load_model_on_init(self):
-        from yolov5 import YOLOv5
-
-        # init model
-        model_path = TestConstants.YOLOV5S_MODEL_PATH
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # check if loaded
-        self.assertNotEqual(yolov5.model, None)
-
-    def test_predict(self):
-        from PIL import Image
-        from yolov5 import YOLOv5
-
-        # init yolov5s model
-        model_path = TestConstants.YOLOV5S_MODEL_PATH
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # prepare image
-        image_path = TestConstants.ZIDANE_IMAGE_PATH
-        image = Image.open(image_path)
-
-        # perform inference
-        results = yolov5.predict(image, size=640, augment=False)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 4)
-
-        # prepare image
-        image = image_path
-
-        # perform inference
-        results = yolov5.predict(image, size=640, augment=True)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 5)
-
-        # init yolov5l model
-        model_path = TestConstants.YOLOV5L_MODEL_PATH
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # prepare image
-        image_path = TestConstants.BUS_IMAGE_PATH
-        image = Image.open(image_path)
-        # perform inference
-        results = yolov5.predict(image, size=1280, augment=False)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 6)
-
-        # prepare image
-        image = image_path
-
-        # perform inference
-        results = yolov5.predict(image, size=1280, augment=False)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 6)
-
-        # init yolov5s model
-        model_path = TestConstants.YOLOV5S_MODEL_PATH
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # prepare images
-        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
-        image_path2 = TestConstants.BUS_IMAGE_PATH
-        image1 = Image.open(image_path1)
-        image2 = Image.open(image_path2)
-
-        # perform inference with multiple images and test augmentation
-        results = yolov5.predict([image1, image2], size=1280, augment=True)
-
-        # compare
-        self.assertEqual(results.n, 2)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 4)
-        self.assertEqual(len(results.pred[1]), 5)
-
-        # prepare image
-        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
-        image_path2 = TestConstants.BUS_IMAGE_PATH
-        image1 = image_path1
-        image2 = image_path2
-
-        # perform inference
-        results = yolov5.predict([image1, image2], size=1280, augment=True)
-
-        # compare
-        self.assertEqual(results.n, 2)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 4)
-        self.assertEqual(len(results.pred[1]), 5)
-
-    def test_hublike_load_model(self):
         import yolov5
 
         # init model
@@ -127,10 +15,11 @@ def test_hublike_load_model(self):
         # check if loaded
         self.assertNotEqual(model, None)
 
-    def test_hublike_predict(self):
-        import yolov5
+    def test_predict(self):
         from PIL import Image
 
+        import yolov5
+
         # init yolov5s model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
         model = yolov5.load(model_path, device=DEVICE)
@@ -182,120 +71,8 @@ def test_hublike_predict(self):
         self.assertEqual(len(results.pred[1]), 5)
 
 
-class TestYolov5FromHuggingface(unittest.TestCase):
+class TestYolov5FromHuggingface(unittest.TestCase):      
     def test_load_model(self):
-        from yolov5 import YOLOv5
-
-        # init model
-        model_path = TestConstants.YOLOV5S_HUB_ID
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=False)
-        yolov5.load_model()
-
-        # check if loaded
-        self.assertNotEqual(yolov5.model, None)
-
-    def test_load_model_on_init(self):
-        from yolov5 import YOLOv5
-
-        # init model
-        model_path = TestConstants.YOLOV5S_HUB_ID
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # check if loaded
-        self.assertNotEqual(yolov5.model, None)
-
-    def test_predict(self):
-        from PIL import Image
-        from yolov5 import YOLOv5
-
-        # init yolov5s model
-        model_path = TestConstants.YOLOV5S_HUB_ID
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # prepare image
-        image_path = TestConstants.ZIDANE_IMAGE_PATH
-        image = Image.open(image_path)
-
-        # perform inference
-        results = yolov5.predict(image, size=640, augment=False)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 4)
-
-        # prepare image
-        image = image_path
-
-        # perform inference
-        results = yolov5.predict(image, size=640, augment=True)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 5)
-
-        # init yolov5l model
-        model_path = TestConstants.YOLOV5L_MODEL_PATH
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # prepare image
-        image_path = TestConstants.BUS_IMAGE_PATH
-        image = Image.open(image_path)
-        # perform inference
-        results = yolov5.predict(image, size=1280, augment=False)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 6)
-
-        # prepare image
-        image = image_path
-
-        # perform inference
-        results = yolov5.predict(image, size=1280, augment=False)
-
-        # compare
-        self.assertEqual(results.n, 1)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 6)
-
-        # init yolov5s model
-        model_path = TestConstants.YOLOV5S_HUB_ID
-        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
-
-        # prepare images
-        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
-        image_path2 = TestConstants.BUS_IMAGE_PATH
-        image1 = Image.open(image_path1)
-        image2 = Image.open(image_path2)
-
-        # perform inference with multiple images and test augmentation
-        results = yolov5.predict([image1, image2], size=1280, augment=True)
-
-        # compare
-        self.assertEqual(results.n, 2)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 4)
-        self.assertEqual(len(results.pred[1]), 5)
-
-        # prepare image
-        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
-        image_path2 = TestConstants.BUS_IMAGE_PATH
-        image1 = image_path1
-        image2 = image_path2
-
-        # perform inference
-        results = yolov5.predict([image1, image2], size=1280, augment=True)
-
-        # compare
-        self.assertEqual(results.n, 2)
-        self.assertEqual(len(results.names), 80)
-        self.assertEqual(len(results.pred[0]), 4)
-        self.assertEqual(len(results.pred[1]), 5)
-
-    def test_hublike_load_model(self):
         import yolov5
 
         # init model
@@ -305,10 +82,11 @@ def test_hublike_load_model(self):
         # check if loaded
         self.assertNotEqual(model, None)
 
-    def test_hublike_predict(self):
-        import yolov5
+    def test_predict(self):
         from PIL import Image
 
+        import yolov5
+
         # init yolov5s model
         model_path = TestConstants.YOLOV5S_HUB_ID
         model = yolov5.load(model_path, device=DEVICE)

yolov5/classify/train.py

@@ -31,6 +31,7 @@
 from tqdm import tqdm
 
 from yolov5.utils.downloads import attempt_donwload_from_hub
+from yolov5.utils.roboflow import check_dataset_roboflow
 
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[1]  # YOLOv5 root directory
@@ -42,12 +43,18 @@
 from yolov5.models.experimental import attempt_load
 from yolov5.models.yolo import ClassificationModel, DetectionModel
 from yolov5.utils.dataloaders import create_classification_dataloader
-from yolov5.utils.general import (DATASETS_DIR, LOGGER, TQDM_BAR_FORMAT, WorkingDirectory, check_git_info, check_git_status,
-                           check_requirements, colorstr, download, increment_path, init_seeds, print_args, yaml_save)
+from yolov5.utils.general import (DATASETS_DIR, LOGGER, TQDM_BAR_FORMAT,
+                                  WorkingDirectory, check_git_info,
+                                  check_git_status, check_requirements,
+                                  colorstr, download, increment_path,
+                                  init_seeds, print_args, yaml_save)
 from yolov5.utils.loggers import GenericLogger
 from yolov5.utils.plots import imshow_cls
-from yolov5.utils.torch_utils import (ModelEMA, model_info, reshape_classifier_output, select_device, smart_DDP,
-                               smart_optimizer, smartCrossEntropyLoss, torch_distributed_zero_first)
+from yolov5.utils.torch_utils import (ModelEMA, model_info,
+                                      reshape_classifier_output, select_device,
+                                      smart_DDP, smart_optimizer,
+                                      smartCrossEntropyLoss,
+                                      torch_distributed_zero_first)
 
 LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
 RANK = int(os.getenv('RANK', -1))
@@ -302,6 +309,9 @@ def parse_opt(known=False):
     parser.add_argument('--neptune_token', type=str, default=None, help='neptune.ai api token')
     parser.add_argument('--neptune_project', type=str, default=None, help='https://docs.neptune.ai/api-reference/neptune')
 
+    # Roboflow arguments
+    parser.add_argument('--roboflow_token', type=str, default=None, help='roboflow api token')
+
     return parser.parse_known_args()[0] if known else parser.parse_args()
 
 
@@ -312,6 +322,14 @@ def main(opt):
         check_git_status()
         check_requirements()
 
+    if "roboflow.com" in str(opt.data):
+        opt.data = check_dataset_roboflow(
+            data=opt.data,
+            roboflow_token=opt.roboflow_token,
+            task="classify",
+            location=ROOT.absolute().as_posix()
+        )
+
     # DDP mode
     device = select_device(opt.device, batch_size=opt.batch_size)
     if LOCAL_RANK != -1: