Model builder (#29)

Co-authored-by: Ayush Chaurasia <ayush.chuararsia@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/data/__init__.py

@@ -1,3 +1,4 @@
+from .base import BaseDataset
 from .build import build_classification_dataloader, build_dataloader
 from .dataset import ClassificationDataset, SemanticDataset, YOLODataset
 from .dataset_wrappers import MixAndRectDataset

ultralytics/yolo/data/augment.py

@@ -8,9 +8,12 @@ import numpy as np
 import torch
 import torchvision.transforms as T
 
-from ..utils.general import LOGGER, check_version, colorstr, segment2box
+from ..utils import LOGGER
+from ..utils.checks import check_version
 from ..utils.instance import Instances
+from ..utils.loggers import colorstr
 from ..utils.metrics import bbox_ioa
+from ..utils.ops import segment2box
 from .utils import IMAGENET_MEAN, IMAGENET_STD, polygons2masks, polygons2masks_overlap
 
 
@@ -457,8 +460,8 @@ class LetterBox:
         self.scaleup = scaleup
         self.stride = stride
 
-    def __call__(self, labels):
-        img = labels["img"]
+    def __call__(self, labels={}, image=None):
+        img = image or labels["img"]
         shape = img.shape[:2]  # current shape [height, width]
         new_shape = labels.get("rect_shape", self.new_shape)
         if isinstance(new_shape, int):

ultralytics/yolo/data/base.py

@@ -9,7 +9,7 @@ import numpy as np
 from torch.utils.data import Dataset
 from tqdm import tqdm
 
-from ..utils.general import NUM_THREADS
+from ..utils import NUM_THREADS
 from .utils import BAR_FORMAT, HELP_URL, IMG_FORMATS, LOCAL_RANK
 
 

ultralytics/yolo/data/build.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 from torch.utils.data import DataLoader, dataloader, distributed
 
-from ..utils.general import LOGGER
+from ..utils import LOGGER
 from ..utils.torch_utils import torch_distributed_zero_first
 from .dataset import ClassificationDataset, YOLODataset
 from .utils import PIN_MEMORY, RANK

ultralytics/yolo/data/dataset.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import torchvision
 from tqdm import tqdm
 
-from ..utils.general import NUM_THREADS
+from ..utils import NUM_THREADS
 from .augment import *
 from .base import BaseDataset
 from .utils import BAR_FORMAT, HELP_URL, LOCAL_RANK, get_hash, img2label_paths, verify_image_label

ultralytics/yolo/data/utils.py

@@ -6,7 +6,7 @@ import cv2
 import numpy as np
 from PIL import ExifTags, Image, ImageOps
 
-from ..utils.general import segments2boxes
+from ..utils.ops import segments2boxes
 
 HELP_URL = "See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data"
 IMG_FORMATS = "bmp", "dng", "jpeg", "jpg", "mpo", "png", "tif", "tiff", "webp", "pfm"  # include image suffixes