ultralytics 8.0.169TQDM, INTERP_LINEAR and RTDETR load_image() updates (#4704)

Co-authored-by: Rustem Galiullin <rustemgal@gmail.com>
Co-authored-by: Rustem Galiullin <rustem.galiullin@bayanat.ai>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-09-02 20:01:57 +02:00 committed by GitHub
parent a4fabfdacf
commit 187b504d68
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 101 additions and 120 deletions

View file

@ -8,9 +8,8 @@ import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm
from ultralytics.utils import TQDM_BAR_FORMAT
from ultralytics.utils import TQDM
class FastSAMPrompt:
@ -87,7 +86,7 @@ class FastSAMPrompt:
retina=False,
withContours=True):
n = len(annotations)
pbar = tqdm(annotations, total=n, bar_format=TQDM_BAR_FORMAT)
pbar = TQDM(annotations, total=n)
for ann in pbar:
result_name = os.path.basename(ann.path)
image = ann.orig_img
@ -156,7 +155,7 @@ class FastSAMPrompt:
save_path.parent.mkdir(exist_ok=True, parents=True)
cv2.imwrite(str(save_path), img_array)
plt.close()
pbar.set_description('Saving {} to {}'.format(result_name, save_path))
pbar.set_description(f'Saving {result_name} to {save_path}')
@staticmethod
def fast_show_mask(

View file

@ -2,8 +2,6 @@
from pathlib import Path
import cv2
import numpy as np
import torch
from ultralytics.data import YOLODataset
@ -21,30 +19,9 @@ class RTDETRDataset(YOLODataset):
super().__init__(*args, data=data, use_segments=False, use_keypoints=False, **kwargs)
# NOTE: add stretch version load_image for rtdetr mosaic
def load_image(self, i):
def load_image(self, i, rect_mode=False):
"""Loads 1 image from dataset index 'i', returns (im, resized hw)."""
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
if im is None: # not cached in RAM
if fn.exists(): # load npy
im = np.load(fn)
else: # read image
im = cv2.imread(f) # BGR
if im is None:
raise FileNotFoundError(f'Image Not Found {f}')
h0, w0 = im.shape[:2] # orig hw
im = cv2.resize(im, (self.imgsz, self.imgsz), interpolation=cv2.INTER_LINEAR)
# Add to buffer if training with augmentations
if self.augment:
self.ims[i], self.im_hw0[i], self.im_hw[i] = im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
self.buffer.append(i)
if len(self.buffer) >= self.max_buffer_length:
j = self.buffer.pop(0)
self.ims[j], self.im_hw0[j], self.im_hw[j] = None, None, None
return im, (h0, w0), im.shape[:2]
return self.ims[i], self.im_hw0[i], self.im_hw[i]
return super().load_image(i=i, rect_mode=rect_mode)
def build_transforms(self, hyp=None):
"""Temporary, only for evaluation."""