ultralytics 8.0.238__py3-none-any.whl → 8.0.239__py3-none-any.whl

ultralytics/data/loaders.py

@@ -23,6 +23,7 @@ from ultralytics.utils.checks import check_requirements
 @dataclass
 class SourceTypes:
     """Class to represent various types of input sources for predictions."""
+
     webcam: bool = False
     screenshot: bool = False
     from_img: bool = False
@@ -59,12 +60,12 @@ class LoadStreams:
         __len__: Return the length of the sources object.
     """
 
-    def __init__(self, sources='file.streams', imgsz=640, vid_stride=1, buffer=False):
+    def __init__(self, sources="file.streams", imgsz=640, vid_stride=1, buffer=False):
         """Initialize instance variables and check for consistent input stream shapes."""
         torch.backends.cudnn.benchmark = True  # faster for fixed-size inference
         self.buffer = buffer  # buffer input streams
         self.running = True  # running flag for Thread
-        self.mode = 'stream'
+        self.mode = "stream"
         self.imgsz = imgsz
         self.vid_stride = vid_stride  # video frame-rate stride
 
@@ -79,33 +80,36 @@ class LoadStreams:
         self.sources = [ops.clean_str(x) for x in sources]  # clean source names for later
         for i, s in enumerate(sources):  # index, source
             # Start thread to read frames from video stream
-            st = f'{i + 1}/{n}: {s}... '
-            if urlparse(s).hostname in ('www.youtube.com', 'youtube.com', 'youtu.be'):  # if source is YouTube video
+            st = f"{i + 1}/{n}: {s}... "
+            if urlparse(s).hostname in ("www.youtube.com", "youtube.com", "youtu.be"):  # if source is YouTube video
                 # YouTube format i.e. 'https://www.youtube.com/watch?v=Zgi9g1ksQHc' or 'https://youtu.be/LNwODJXcvt4'
                 s = get_best_youtube_url(s)
             s = eval(s) if s.isnumeric() else s  # i.e. s = '0' local webcam
             if s == 0 and (is_colab() or is_kaggle()):
-                raise NotImplementedError("'source=0' webcam not supported in Colab and Kaggle notebooks. "
-                                          "Try running 'source=0' in a local environment.")
+                raise NotImplementedError(
+                    "'source=0' webcam not supported in Colab and Kaggle notebooks. "
+                    "Try running 'source=0' in a local environment."
+                )
             self.caps[i] = cv2.VideoCapture(s)  # store video capture object
             if not self.caps[i].isOpened():
-                raise ConnectionError(f'{st}Failed to open {s}')
+                raise ConnectionError(f"{st}Failed to open {s}")
             w = int(self.caps[i].get(cv2.CAP_PROP_FRAME_WIDTH))
             h = int(self.caps[i].get(cv2.CAP_PROP_FRAME_HEIGHT))
             fps = self.caps[i].get(cv2.CAP_PROP_FPS)  # warning: may return 0 or nan
             self.frames[i] = max(int(self.caps[i].get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float(
-                'inf')  # infinite stream fallback
+                "inf"
+            )  # infinite stream fallback
             self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30  # 30 FPS fallback
 
             success, im = self.caps[i].read()  # guarantee first frame
             if not success or im is None:
-                raise ConnectionError(f'{st}Failed to read images from {s}')
+                raise ConnectionError(f"{st}Failed to read images from {s}")
             self.imgs[i].append(im)
             self.shape[i] = im.shape
             self.threads[i] = Thread(target=self.update, args=([i, self.caps[i], s]), daemon=True)
-            LOGGER.info(f'{st}Success ✅ ({self.frames[i]} frames of shape {w}x{h} at {self.fps[i]:.2f} FPS)')
+            LOGGER.info(f"{st}Success ✅ ({self.frames[i]} frames of shape {w}x{h} at {self.fps[i]:.2f} FPS)")
             self.threads[i].start()
-        LOGGER.info('')  # newline
+        LOGGER.info("")  # newline
 
         # Check for common shapes
         self.bs = self.__len__()
@@ -121,7 +125,7 @@ class LoadStreams:
                     success, im = cap.retrieve()
                     if not success:
                         im = np.zeros(self.shape[i], dtype=np.uint8)
-                        LOGGER.warning('WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.')
+                        LOGGER.warning("WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.")
                         cap.open(stream)  # re-open stream if signal was lost
                     if self.buffer:
                         self.imgs[i].append(im)
@@ -140,7 +144,7 @@ class LoadStreams:
             try:
                 cap.release()  # release video capture
             except Exception as e:
-                LOGGER.warning(f'WARNING ⚠️ Could not release VideoCapture object: {e}')
+                LOGGER.warning(f"WARNING ⚠️ Could not release VideoCapture object: {e}")
         cv2.destroyAllWindows()
 
     def __iter__(self):
@@ -154,16 +158,15 @@ class LoadStreams:
 
         images = []
         for i, x in enumerate(self.imgs):
-
             # Wait until a frame is available in each buffer
             while not x:
-                if not self.threads[i].is_alive() or cv2.waitKey(1) == ord('q'):  # q to quit
+                if not self.threads[i].is_alive() or cv2.waitKey(1) == ord("q"):  # q to quit
                     self.close()
                     raise StopIteration
                 time.sleep(1 / min(self.fps))
                 x = self.imgs[i]
                 if not x:
-                    LOGGER.warning(f'WARNING ⚠️ Waiting for stream {i}')
+                    LOGGER.warning(f"WARNING ⚠️ Waiting for stream {i}")
 
             # Get and remove the first frame from imgs buffer
             if self.buffer:
@@ -174,7 +177,7 @@ class LoadStreams:
                 images.append(x.pop(-1) if x else np.zeros(self.shape[i], dtype=np.uint8))
                 x.clear()
 
-        return self.sources, images, None, ''
+        return self.sources, images, None, ""
 
     def __len__(self):
         """Return the length of the sources object."""
@@ -209,7 +212,7 @@ class LoadScreenshots:
 
     def __init__(self, source, imgsz=640):
         """Source = [screen_number left top width height] (pixels)."""
-        check_requirements('mss')
+        check_requirements("mss")
         import mss  # noqa
 
         source, *params = source.split()
@@ -221,18 +224,18 @@ class LoadScreenshots:
         elif len(params) == 5:
             self.screen, left, top, width, height = (int(x) for x in params)
         self.imgsz = imgsz
-        self.mode = 'stream'
+        self.mode = "stream"
         self.frame = 0
         self.sct = mss.mss()
         self.bs = 1
 
         # Parse monitor shape
         monitor = self.sct.monitors[self.screen]
-        self.top = monitor['top'] if top is None else (monitor['top'] + top)
-        self.left = monitor['left'] if left is None else (monitor['left'] + left)
-        self.width = width or monitor['width']
-        self.height = height or monitor['height']
-        self.monitor = {'left': self.left, 'top': self.top, 'width': self.width, 'height': self.height}
+        self.top = monitor["top"] if top is None else (monitor["top"] + top)
+        self.left = monitor["left"] if left is None else (monitor["left"] + left)
+        self.width = width or monitor["width"]
+        self.height = height or monitor["height"]
+        self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}
 
     def __iter__(self):
         """Returns an iterator of the object."""
@@ -241,7 +244,7 @@ class LoadScreenshots:
     def __next__(self):
         """mss screen capture: get raw pixels from the screen as np array."""
         im0 = np.asarray(self.sct.grab(self.monitor))[:, :, :3]  # BGRA to BGR
-        s = f'screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: '
+        s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: "
 
         self.frame += 1
         return [str(self.screen)], [im0], None, s  # screen, img, vid_cap, string
@@ -274,32 +277,32 @@ class LoadImages:
     def __init__(self, path, imgsz=640, vid_stride=1):
         """Initialize the Dataloader and raise FileNotFoundError if file not found."""
         parent = None
-        if isinstance(path, str) and Path(path).suffix == '.txt':  # *.txt file with img/vid/dir on each line
+        if isinstance(path, str) and Path(path).suffix == ".txt":  # *.txt file with img/vid/dir on each line
             parent = Path(path).parent
             path = Path(path).read_text().splitlines()  # list of sources
         files = []
         for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
             a = str(Path(p).absolute())  # do not use .resolve() https://github.com/ultralytics/ultralytics/issues/2912
-            if '*' in a:
+            if "*" in a:
                 files.extend(sorted(glob.glob(a, recursive=True)))  # glob
             elif os.path.isdir(a):
-                files.extend(sorted(glob.glob(os.path.join(a, '*.*'))))  # dir
+                files.extend(sorted(glob.glob(os.path.join(a, "*.*"))))  # dir
             elif os.path.isfile(a):
                 files.append(a)  # files (absolute or relative to CWD)
             elif parent and (parent / p).is_file():
                 files.append(str((parent / p).absolute()))  # files (relative to *.txt file parent)
             else:
-                raise FileNotFoundError(f'{p} does not exist')
+                raise FileNotFoundError(f"{p} does not exist")
 
-        images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
-        videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
+        images = [x for x in files if x.split(".")[-1].lower() in IMG_FORMATS]
+        videos = [x for x in files if x.split(".")[-1].lower() in VID_FORMATS]
         ni, nv = len(images), len(videos)
 
         self.imgsz = imgsz
         self.files = images + videos
         self.nf = ni + nv  # number of files
         self.video_flag = [False] * ni + [True] * nv
-        self.mode = 'image'
+        self.mode = "image"
         self.vid_stride = vid_stride  # video frame-rate stride
         self.bs = 1
         if any(videos):
@@ -307,8 +310,10 @@ class LoadImages:
         else:
             self.cap = None
         if self.nf == 0:
-            raise FileNotFoundError(f'No images or videos found in {p}. '
-                                    f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}')
+            raise FileNotFoundError(
+                f"No images or videos found in {p}. "
+                f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
+            )
 
     def __iter__(self):
         """Returns an iterator object for VideoStream or ImageFolder."""
@@ -323,7 +328,7 @@ class LoadImages:
 
         if self.video_flag[self.count]:
             # Read video
-            self.mode = 'video'
+            self.mode = "video"
             for _ in range(self.vid_stride):
                 self.cap.grab()
             success, im0 = self.cap.retrieve()
@@ -338,15 +343,15 @@ class LoadImages:
 
             self.frame += 1
             # im0 = self._cv2_rotate(im0)  # for use if cv2 autorotation is False
-            s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '
+            s = f"video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: "
 
         else:
             # Read image
             self.count += 1
             im0 = cv2.imread(path)  # BGR
             if im0 is None:
-                raise FileNotFoundError(f'Image Not Found {path}')
-            s = f'image {self.count}/{self.nf} {path}: '
+                raise FileNotFoundError(f"Image Not Found {path}")
+            s = f"image {self.count}/{self.nf} {path}: "
 
         return [path], [im0], self.cap, s
 
@@ -385,20 +390,20 @@ class LoadPilAndNumpy:
         """Initialize PIL and Numpy Dataloader."""
         if not isinstance(im0, list):
             im0 = [im0]
-        self.paths = [getattr(im, 'filename', f'image{i}.jpg') for i, im in enumerate(im0)]
+        self.paths = [getattr(im, "filename", f"image{i}.jpg") for i, im in enumerate(im0)]
         self.im0 = [self._single_check(im) for im in im0]
         self.imgsz = imgsz
-        self.mode = 'image'
+        self.mode = "image"
         # Generate fake paths
         self.bs = len(self.im0)
 
     @staticmethod
     def _single_check(im):
         """Validate and format an image to numpy array."""
-        assert isinstance(im, (Image.Image, np.ndarray)), f'Expected PIL/np.ndarray image type, but got {type(im)}'
+        assert isinstance(im, (Image.Image, np.ndarray)), f"Expected PIL/np.ndarray image type, but got {type(im)}"
         if isinstance(im, Image.Image):
-            if im.mode != 'RGB':
-                im = im.convert('RGB')
+            if im.mode != "RGB":
+                im = im.convert("RGB")
             im = np.asarray(im)[:, :, ::-1]
             im = np.ascontiguousarray(im)  # contiguous
         return im
@@ -412,7 +417,7 @@ class LoadPilAndNumpy:
         if self.count == 1:  # loop only once as it's batch inference
             raise StopIteration
         self.count += 1
-        return self.paths, self.im0, None, ''
+        return self.paths, self.im0, None, ""
 
     def __iter__(self):
         """Enables iteration for class LoadPilAndNumpy."""
@@ -441,14 +446,16 @@ class LoadTensor:
         """Initialize Tensor Dataloader."""
         self.im0 = self._single_check(im0)
         self.bs = self.im0.shape[0]
-        self.mode = 'image'
-        self.paths = [getattr(im, 'filename', f'image{i}.jpg') for i, im in enumerate(im0)]
+        self.mode = "image"
+        self.paths = [getattr(im, "filename", f"image{i}.jpg") for i, im in enumerate(im0)]
 
     @staticmethod
     def _single_check(im, stride=32):
         """Validate and format an image to torch.Tensor."""
-        s = f'WARNING ⚠️ torch.Tensor inputs should be BCHW i.e. shape(1, 3, 640, 640) ' \
-            f'divisible by stride {stride}. Input shape{tuple(im.shape)} is incompatible.'
+        s = (
+            f"WARNING ⚠️ torch.Tensor inputs should be BCHW i.e. shape(1, 3, 640, 640) "
+            f"divisible by stride {stride}. Input shape{tuple(im.shape)} is incompatible."
+        )
         if len(im.shape) != 4:
             if len(im.shape) != 3:
                 raise ValueError(s)
@@ -457,8 +464,10 @@ class LoadTensor:
         if im.shape[2] % stride or im.shape[3] % stride:
             raise ValueError(s)
         if im.max() > 1.0 + torch.finfo(im.dtype).eps:  # torch.float32 eps is 1.2e-07
-            LOGGER.warning(f'WARNING ⚠️ torch.Tensor inputs should be normalized 0.0-1.0 but max value is {im.max()}. '
-                           f'Dividing input by 255.')
+            LOGGER.warning(
+                f"WARNING ⚠️ torch.Tensor inputs should be normalized 0.0-1.0 but max value is {im.max()}. "
+                f"Dividing input by 255."
+            )
             im = im.float() / 255.0
 
         return im
@@ -473,7 +482,7 @@ class LoadTensor:
         if self.count == 1:
             raise StopIteration
         self.count += 1
-        return self.paths, self.im0, None, ''
+        return self.paths, self.im0, None, ""
 
     def __len__(self):
         """Returns the batch size."""
@@ -485,12 +494,14 @@ def autocast_list(source):
     files = []
     for im in source:
         if isinstance(im, (str, Path)):  # filename or uri
-            files.append(Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im))
+            files.append(Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im))
         elif isinstance(im, (Image.Image, np.ndarray)):  # PIL or np Image
             files.append(im)
         else:
-            raise TypeError(f'type {type(im).__name__} is not a supported Ultralytics prediction source type. \n'
-                            f'See https://docs.ultralytics.com/modes/predict for supported source types.')
+            raise TypeError(
+                f"type {type(im).__name__} is not a supported Ultralytics prediction source type. \n"
+                f"See https://docs.ultralytics.com/modes/predict for supported source types."
+            )
 
     return files
 
@@ -513,16 +524,18 @@ def get_best_youtube_url(url, use_pafy=True):
         (str): The URL of the best quality MP4 video stream, or None if no suitable stream is found.
     """
     if use_pafy:
-        check_requirements(('pafy', 'youtube_dl==2020.12.2'))
+        check_requirements(("pafy", "youtube_dl==2020.12.2"))
         import pafy  # noqa
-        return pafy.new(url).getbestvideo(preftype='mp4').url
+
+        return pafy.new(url).getbestvideo(preftype="mp4").url
     else:
-        check_requirements('yt-dlp')
+        check_requirements("yt-dlp")
         import yt_dlp
-        with yt_dlp.YoutubeDL({'quiet': True}) as ydl:
+
+        with yt_dlp.YoutubeDL({"quiet": True}) as ydl:
             info_dict = ydl.extract_info(url, download=False)  # extract info
-        for f in reversed(info_dict.get('formats', [])):  # reversed because best is usually last
+        for f in reversed(info_dict.get("formats", [])):  # reversed because best is usually last
             # Find a format with video codec, no audio, *.mp4 extension at least 1920x1080 size
-            good_size = (f.get('width') or 0) >= 1920 or (f.get('height') or 0) >= 1080
-            if good_size and f['vcodec'] != 'none' and f['acodec'] == 'none' and f['ext'] == 'mp4':
-                return f.get('url')
+            good_size = (f.get("width") or 0) >= 1920 or (f.get("height") or 0) >= 1080
+            if good_size and f["vcodec"] != "none" and f["acodec"] == "none" and f["ext"] == "mp4":
+                return f.get("url")

ultralytics/data/split_dota.py

@@ -1,3 +1,5 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
 import itertools
 import os
 from glob import glob
@@ -12,7 +14,7 @@ from tqdm import tqdm
 from ultralytics.data.utils import exif_size, img2label_paths
 from ultralytics.utils.checks import check_requirements
 
-check_requirements('shapely')
+check_requirements("shapely")
 from shapely.geometry import Polygon
 
 
@@ -52,11 +54,14 @@ def bbox_iof(polygon1, bbox2, eps=1e-6):
     return outputs
 
 
-def load_yolo_dota(data_root, split='train'):
-    """Load DOTA dataset.
+def load_yolo_dota(data_root, split="train"):
+    """
+    Load DOTA dataset.
+
     Args:
         data_root (str): Data root.
         split (str): The split data set, could be train or val.
+
     Notes:
         The directory structure assumed for the DOTA dataset:
             - data_root
@@ -67,10 +72,10 @@ def load_yolo_dota(data_root, split='train'):
                     - train
                     - val
     """
-    assert split in ['train', 'val']
-    im_dir = os.path.join(data_root, f'images/{split}')
+    assert split in ["train", "val"]
+    im_dir = os.path.join(data_root, f"images/{split}")
     assert Path(im_dir).exists(), f"Can't find {im_dir}, please check your data root."
-    im_files = glob(os.path.join(data_root, f'images/{split}/*'))
+    im_files = glob(os.path.join(data_root, f"images/{split}/*"))
     lb_files = img2label_paths(im_files)
     annos = []
     for im_file, lb_file in zip(im_files, lb_files):
@@ -95,7 +100,7 @@ def get_windows(im_size, crop_sizes=[1024], gaps=[200], im_rate_thr=0.6, eps=0.0
     h, w = im_size
     windows = []
     for crop_size, gap in zip(crop_sizes, gaps):
-        assert crop_size > gap, f'invaild crop_size gap pair [{crop_size} {gap}]'
+        assert crop_size > gap, f"invalid crop_size gap pair [{crop_size} {gap}]"
         step = crop_size - gap
 
         xn = 1 if w <= crop_size else ceil((w - crop_size) / step + 1)
@@ -127,13 +132,13 @@ def get_windows(im_size, crop_sizes=[1024], gaps=[200], im_rate_thr=0.6, eps=0.0
 
 def get_window_obj(anno, windows, iof_thr=0.7):
     """Get objects for each window."""
-    h, w = anno['ori_size']
-    label = anno['label']
+    h, w = anno["ori_size"]
+    label = anno["label"]
     if len(label):
         label[:, 1::2] *= w
         label[:, 2::2] *= h
         iofs = bbox_iof(label[:, 1:], windows)
-        # unnormalized and misaligned coordinates
+        # Unnormalized and misaligned coordinates
         window_anns = [(label[iofs[:, i] >= iof_thr]) for i in range(len(windows))]
     else:
         window_anns = [np.zeros((0, 9), dtype=np.float32) for _ in range(len(windows))]
@@ -141,13 +146,16 @@ def get_window_obj(anno, windows, iof_thr=0.7):
 
 
 def crop_and_save(anno, windows, window_objs, im_dir, lb_dir):
-    """Crop images and save new labels.
+    """
+    Crop images and save new labels.
+
     Args:
         anno (dict): Annotation dict, including `filepath`, `label`, `ori_size` as its keys.
         windows (list): A list of windows coordinates.
         window_objs (list): A list of labels inside each window.
         im_dir (str): The output directory path of images.
         lb_dir (str): The output directory path of labels.
+
     Notes:
         The directory structure assumed for the DOTA dataset:
             - data_root
@@ -158,15 +166,15 @@ def crop_and_save(anno, windows, window_objs, im_dir, lb_dir):
                     - train
                     - val
     """
-    im = cv2.imread(anno['filepath'])
-    name = Path(anno['filepath']).stem
+    im = cv2.imread(anno["filepath"])
+    name = Path(anno["filepath"]).stem
     for i, window in enumerate(windows):
         x_start, y_start, x_stop, y_stop = window.tolist()
-        new_name = name + '__' + str(x_stop - x_start) + '__' + str(x_start) + '___' + str(y_start)
+        new_name = name + "__" + str(x_stop - x_start) + "__" + str(x_start) + "___" + str(y_start)
         patch_im = im[y_start:y_stop, x_start:x_stop]
         ph, pw = patch_im.shape[:2]
 
-        cv2.imwrite(os.path.join(im_dir, f'{new_name}.jpg'), patch_im)
+        cv2.imwrite(os.path.join(im_dir, f"{new_name}.jpg"), patch_im)
         label = window_objs[i]
         if len(label) == 0:
             continue
@@ -175,17 +183,17 @@ def crop_and_save(anno, windows, window_objs, im_dir, lb_dir):
         label[:, 1::2] /= pw
         label[:, 2::2] /= ph
 
-        with open(os.path.join(lb_dir, f'{new_name}.txt'), 'w') as f:
+        with open(os.path.join(lb_dir, f"{new_name}.txt"), "w") as f:
             for lb in label:
-                formatted_coords = ['{:.6g}'.format(coord) for coord in lb[1:]]
+                formatted_coords = ["{:.6g}".format(coord) for coord in lb[1:]]
                 f.write(f"{int(lb[0])} {' '.join(formatted_coords)}\n")
 
 
-def split_images_and_labels(data_root, save_dir, split='train', crop_sizes=[1024], gaps=[200]):
+def split_images_and_labels(data_root, save_dir, split="train", crop_sizes=[1024], gaps=[200]):
     """
     Split both images and labels.
 
-    NOTES:
+    Notes:
         The directory structure assumed for the DOTA dataset:
             - data_root
                 - images
@@ -199,14 +207,14 @@ def split_images_and_labels(data_root, save_dir, split='train', crop_sizes=[1024
                 - labels
                     - split
     """
-    im_dir = Path(save_dir) / 'images' / split
+    im_dir = Path(save_dir) / "images" / split
     im_dir.mkdir(parents=True, exist_ok=True)
-    lb_dir = Path(save_dir) / 'labels' / split
+    lb_dir = Path(save_dir) / "labels" / split
     lb_dir.mkdir(parents=True, exist_ok=True)
 
     annos = load_yolo_dota(data_root, split=split)
     for anno in tqdm(annos, total=len(annos), desc=split):
-        windows = get_windows(anno['ori_size'], crop_sizes, gaps)
+        windows = get_windows(anno["ori_size"], crop_sizes, gaps)
         window_objs = get_window_obj(anno, windows)
         crop_and_save(anno, windows, window_objs, str(im_dir), str(lb_dir))
 
@@ -215,7 +223,7 @@ def split_trainval(data_root, save_dir, crop_size=1024, gap=200, rates=[1.0]):
     """
     Split train and val set of DOTA.
 
-    NOTES:
+    Notes:
         The directory structure assumed for the DOTA dataset:
             - data_root
                 - images
@@ -237,7 +245,7 @@ def split_trainval(data_root, save_dir, crop_size=1024, gap=200, rates=[1.0]):
     for r in rates:
         crop_sizes.append(int(crop_size / r))
         gaps.append(int(gap / r))
-    for split in ['train', 'val']:
+    for split in ["train", "val"]:
         split_images_and_labels(data_root, save_dir, split, crop_sizes, gaps)
 
 
@@ -245,7 +253,7 @@ def split_test(data_root, save_dir, crop_size=1024, gap=200, rates=[1.0]):
     """
     Split test set of DOTA, labels are not included within this set.
 
-    NOTES:
+    Notes:
         The directory structure assumed for the DOTA dataset:
             - data_root
                 - images
@@ -259,30 +267,30 @@ def split_test(data_root, save_dir, crop_size=1024, gap=200, rates=[1.0]):
     for r in rates:
         crop_sizes.append(int(crop_size / r))
         gaps.append(int(gap / r))
-    save_dir = Path(save_dir) / 'images' / 'test'
+    save_dir = Path(save_dir) / "images" / "test"
     save_dir.mkdir(parents=True, exist_ok=True)
 
-    im_dir = Path(os.path.join(data_root, 'images/test'))
+    im_dir = Path(os.path.join(data_root, "images/test"))
     assert im_dir.exists(), f"Can't find {str(im_dir)}, please check your data root."
-    im_files = glob(str(im_dir / '*'))
-    for im_file in tqdm(im_files, total=len(im_files), desc='test'):
+    im_files = glob(str(im_dir / "*"))
+    for im_file in tqdm(im_files, total=len(im_files), desc="test"):
         w, h = exif_size(Image.open(im_file))
         windows = get_windows((h, w), crop_sizes=crop_sizes, gaps=gaps)
         im = cv2.imread(im_file)
         name = Path(im_file).stem
         for window in windows:
             x_start, y_start, x_stop, y_stop = window.tolist()
-            new_name = (name + '__' + str(x_stop - x_start) + '__' + str(x_start) + '___' + str(y_start))
+            new_name = name + "__" + str(x_stop - x_start) + "__" + str(x_start) + "___" + str(y_start)
             patch_im = im[y_start:y_stop, x_start:x_stop]
-            cv2.imwrite(os.path.join(str(save_dir), f'{new_name}.jpg'), patch_im)
+            cv2.imwrite(os.path.join(str(save_dir), f"{new_name}.jpg"), patch_im)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     split_trainval(
-        data_root='DOTAv2',
-        save_dir='DOTAv2-split',
+        data_root="DOTAv2",
+        save_dir="DOTAv2-split",
     )
     split_test(
-        data_root='DOTAv2',
-        save_dir='DOTAv2-split',
+        data_root="DOTAv2",
+        save_dir="DOTAv2-split",
     )