diff --git a/ultralytics/data/split_dota.py b/ultralytics/data/split_dota.py index a8f17d41..1d1ca310 100644 --- a/ultralytics/data/split_dota.py +++ b/ultralytics/data/split_dota.py @@ -8,9 +8,9 @@ from pathlib import Path import cv2 import numpy as np from PIL import Image -from tqdm import tqdm from ultralytics.data.utils import exif_size, img2label_paths +from ultralytics.utils import TQDM from ultralytics.utils.checks import check_requirements @@ -221,7 +221,7 @@ def split_images_and_labels(data_root, save_dir, split="train", crop_sizes=(1024 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): + for anno in TQDM(annos, total=len(annos), desc=split): 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)) @@ -281,7 +281,7 @@ def split_test(data_root, save_dir, crop_size=1024, gap=200, rates=(1.0,)): im_dir = Path(data_root) / "images" / "test" assert im_dir.exists(), f"Can't find {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"): + 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) diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py index 4ff46044..f3aee7b5 100644 --- a/ultralytics/utils/__init__.py +++ b/ultralytics/utils/__init__.py @@ -23,8 +23,8 @@ import cv2 import matplotlib.pyplot as plt import numpy as np import torch +import tqdm import yaml -from tqdm import tqdm as tqdm_original from ultralytics import __version__ @@ -133,7 +133,7 @@ os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR" # suppress "NNPACK.cpp could not in os.environ["KINETO_LOG_LEVEL"] = "5" # suppress verbose PyTorch profiler output when computing FLOPs -class TQDM(tqdm_original): +class TQDM(tqdm.tqdm): """ A custom TQDM progress bar class that extends the original tqdm functionality.