add yolo v10 and modify pipeline
This commit is contained in:
@ -22,76 +22,114 @@ from ultralytics.utils.checks import check_requirements
|
||||
|
||||
@dataclass
|
||||
class SourceTypes:
|
||||
webcam: bool = False
|
||||
"""Class to represent various types of input sources for predictions."""
|
||||
|
||||
stream: bool = False
|
||||
screenshot: bool = False
|
||||
from_img: bool = False
|
||||
tensor: bool = False
|
||||
|
||||
|
||||
class LoadStreams:
|
||||
"""YOLOv8 streamloader, i.e. `yolo predict source='rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams`."""
|
||||
"""
|
||||
Stream Loader for various types of video streams, Supports RTSP, RTMP, HTTP, and TCP streams.
|
||||
|
||||
def __init__(self, sources='file.streams', imgsz=640, vid_stride=1, stream_buffer=False):
|
||||
Attributes:
|
||||
sources (str): The source input paths or URLs for the video streams.
|
||||
vid_stride (int): Video frame-rate stride, defaults to 1.
|
||||
buffer (bool): Whether to buffer input streams, defaults to False.
|
||||
running (bool): Flag to indicate if the streaming thread is running.
|
||||
mode (str): Set to 'stream' indicating real-time capture.
|
||||
imgs (list): List of image frames for each stream.
|
||||
fps (list): List of FPS for each stream.
|
||||
frames (list): List of total frames for each stream.
|
||||
threads (list): List of threads for each stream.
|
||||
shape (list): List of shapes for each stream.
|
||||
caps (list): List of cv2.VideoCapture objects for each stream.
|
||||
bs (int): Batch size for processing.
|
||||
|
||||
Methods:
|
||||
__init__: Initialize the stream loader.
|
||||
update: Read stream frames in daemon thread.
|
||||
close: Close stream loader and release resources.
|
||||
__iter__: Returns an iterator object for the class.
|
||||
__next__: Returns source paths, transformed, and original images for processing.
|
||||
__len__: Return the length of the sources object.
|
||||
|
||||
Example:
|
||||
```bash
|
||||
yolo predict source='rtsp://example.com/media.mp4'
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, sources="file.streams", 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.stream_buffer = stream_buffer # buffer input streams
|
||||
self.buffer = buffer # buffer input streams
|
||||
self.running = True # running flag for Thread
|
||||
self.mode = 'stream'
|
||||
self.imgsz = imgsz
|
||||
self.mode = "stream"
|
||||
self.vid_stride = vid_stride # video frame-rate stride
|
||||
|
||||
sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
|
||||
n = len(sources)
|
||||
self.sources = [ops.clean_str(x) for x in sources] # clean source names for later
|
||||
self.imgs, self.fps, self.frames, self.threads, self.shape = [[]] * n, [0] * n, [0] * n, [None] * n, [None] * n
|
||||
self.bs = n
|
||||
self.fps = [0] * n # frames per second
|
||||
self.frames = [0] * n
|
||||
self.threads = [None] * n
|
||||
self.caps = [None] * n # video capture objects
|
||||
self.imgs = [[] for _ in range(n)] # images
|
||||
self.shape = [[] for _ in range(n)] # image shapes
|
||||
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
|
||||
# YouTube format i.e. 'https://www.youtube.com/watch?v=Zgi9g1ksQHc' or 'https://youtu.be/Zgi9g1ksQHc'
|
||||
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
|
||||
|
||||
# Check for common shapes
|
||||
self.bs = self.__len__()
|
||||
LOGGER.info("") # newline
|
||||
|
||||
def update(self, i, cap, stream):
|
||||
"""Read stream `i` frames in daemon thread."""
|
||||
n, f = 0, self.frames[i] # frame number, frame array
|
||||
while self.running and cap.isOpened() and n < (f - 1):
|
||||
# Only read a new frame if the buffer is empty
|
||||
if not self.imgs[i] or not self.stream_buffer:
|
||||
if len(self.imgs[i]) < 30: # keep a <=30-image buffer
|
||||
n += 1
|
||||
cap.grab() # .read() = .grab() followed by .retrieve()
|
||||
if n % self.vid_stride == 0:
|
||||
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
|
||||
self.imgs[i].append(im) # add image to buffer
|
||||
if self.buffer:
|
||||
self.imgs[i].append(im)
|
||||
else:
|
||||
self.imgs[i] = [im]
|
||||
else:
|
||||
time.sleep(0.01) # wait until the buffer is empty
|
||||
|
||||
@ -105,7 +143,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):
|
||||
@ -117,36 +155,62 @@ class LoadStreams:
|
||||
"""Returns source paths, transformed and original images for processing."""
|
||||
self.count += 1
|
||||
|
||||
# Wait until a frame is available in each buffer
|
||||
while not all(self.imgs):
|
||||
if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit
|
||||
self.close()
|
||||
raise StopIteration
|
||||
time.sleep(1 / min(self.fps))
|
||||
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
|
||||
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}")
|
||||
|
||||
# Get and remove the next frame from imgs buffer
|
||||
if self.stream_buffer:
|
||||
images = [x.pop(0) for x in self.imgs]
|
||||
else:
|
||||
# Get the latest frame, and clear the rest from the imgs buffer
|
||||
images = []
|
||||
for x in self.imgs:
|
||||
images.append(x.pop(-1) if x else None)
|
||||
# Get and remove the first frame from imgs buffer
|
||||
if self.buffer:
|
||||
images.append(x.pop(0))
|
||||
|
||||
# Get the last frame, and clear the rest from the imgs buffer
|
||||
else:
|
||||
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, [""] * self.bs
|
||||
|
||||
def __len__(self):
|
||||
"""Return the length of the sources object."""
|
||||
return len(self.sources) # 1E12 frames = 32 streams at 30 FPS for 30 years
|
||||
return self.bs # 1E12 frames = 32 streams at 30 FPS for 30 years
|
||||
|
||||
|
||||
class LoadScreenshots:
|
||||
"""YOLOv8 screenshot dataloader, i.e. `yolo predict source=screen`."""
|
||||
"""
|
||||
YOLOv8 screenshot dataloader.
|
||||
|
||||
def __init__(self, source, imgsz=640):
|
||||
"""source = [screen_number left top width height] (pixels)."""
|
||||
check_requirements('mss')
|
||||
This class manages the loading of screenshot images for processing with YOLOv8.
|
||||
Suitable for use with `yolo predict source=screen`.
|
||||
|
||||
Attributes:
|
||||
source (str): The source input indicating which screen to capture.
|
||||
screen (int): The screen number to capture.
|
||||
left (int): The left coordinate for screen capture area.
|
||||
top (int): The top coordinate for screen capture area.
|
||||
width (int): The width of the screen capture area.
|
||||
height (int): The height of the screen capture area.
|
||||
mode (str): Set to 'stream' indicating real-time capture.
|
||||
frame (int): Counter for captured frames.
|
||||
sct (mss.mss): Screen capture object from `mss` library.
|
||||
bs (int): Batch size, set to 1.
|
||||
monitor (dict): Monitor configuration details.
|
||||
|
||||
Methods:
|
||||
__iter__: Returns an iterator object.
|
||||
__next__: Captures the next screenshot and returns it.
|
||||
"""
|
||||
|
||||
def __init__(self, source):
|
||||
"""Source = [screen_number left top width height] (pixels)."""
|
||||
check_requirements("mss")
|
||||
import mss # noqa
|
||||
|
||||
source, *params = source.split()
|
||||
@ -157,19 +221,19 @@ class LoadScreenshots:
|
||||
left, top, width, height = (int(x) for x in params)
|
||||
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
|
||||
self.fps = 30
|
||||
|
||||
# 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."""
|
||||
@ -178,53 +242,75 @@ 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
|
||||
return [str(self.screen)], [im0], [s] # screen, img, string
|
||||
|
||||
|
||||
class LoadImages:
|
||||
"""YOLOv8 image/video dataloader, i.e. `yolo predict source=image.jpg/vid.mp4`."""
|
||||
class LoadImagesAndVideos:
|
||||
"""
|
||||
YOLOv8 image/video dataloader.
|
||||
|
||||
def __init__(self, path, imgsz=640, vid_stride=1):
|
||||
This class manages the loading and pre-processing of image and video data for YOLOv8. It supports loading from
|
||||
various formats, including single image files, video files, and lists of image and video paths.
|
||||
|
||||
Attributes:
|
||||
files (list): List of image and video file paths.
|
||||
nf (int): Total number of files (images and videos).
|
||||
video_flag (list): Flags indicating whether a file is a video (True) or an image (False).
|
||||
mode (str): Current mode, 'image' or 'video'.
|
||||
vid_stride (int): Stride for video frame-rate, defaults to 1.
|
||||
bs (int): Batch size, set to 1 for this class.
|
||||
cap (cv2.VideoCapture): Video capture object for OpenCV.
|
||||
frame (int): Frame counter for video.
|
||||
frames (int): Total number of frames in the video.
|
||||
count (int): Counter for iteration, initialized at 0 during `__iter__()`.
|
||||
|
||||
Methods:
|
||||
_new_video(path): Create a new cv2.VideoCapture object for a given video path.
|
||||
"""
|
||||
|
||||
def __init__(self, path, batch=1, 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.ni = ni # number of images
|
||||
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
|
||||
self.bs = batch
|
||||
if any(videos):
|
||||
self._new_video(videos[0]) # new video
|
||||
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."""
|
||||
@ -232,71 +318,105 @@ class LoadImages:
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
"""Return next image, path and metadata from dataset."""
|
||||
if self.count == self.nf:
|
||||
raise StopIteration
|
||||
path = self.files[self.count]
|
||||
|
||||
if self.video_flag[self.count]:
|
||||
# Read video
|
||||
self.mode = 'video'
|
||||
for _ in range(self.vid_stride):
|
||||
self.cap.grab()
|
||||
success, im0 = self.cap.retrieve()
|
||||
while not success:
|
||||
self.count += 1
|
||||
self.cap.release()
|
||||
if self.count == self.nf: # last video
|
||||
"""Returns the next batch of images or video frames along with their paths and metadata."""
|
||||
paths, imgs, info = [], [], []
|
||||
while len(imgs) < self.bs:
|
||||
if self.count >= self.nf: # end of file list
|
||||
if len(imgs) > 0:
|
||||
return paths, imgs, info # return last partial batch
|
||||
else:
|
||||
raise StopIteration
|
||||
path = self.files[self.count]
|
||||
self._new_video(path)
|
||||
success, im0 = self.cap.read()
|
||||
|
||||
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}: '
|
||||
path = self.files[self.count]
|
||||
if self.video_flag[self.count]:
|
||||
self.mode = "video"
|
||||
if not self.cap or not self.cap.isOpened():
|
||||
self._new_video(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}: '
|
||||
for _ in range(self.vid_stride):
|
||||
success = self.cap.grab()
|
||||
if not success:
|
||||
break # end of video or failure
|
||||
|
||||
return [path], [im0], self.cap, s
|
||||
if success:
|
||||
success, im0 = self.cap.retrieve()
|
||||
if success:
|
||||
self.frame += 1
|
||||
paths.append(path)
|
||||
imgs.append(im0)
|
||||
info.append(f"video {self.count + 1}/{self.nf} (frame {self.frame}/{self.frames}) {path}: ")
|
||||
if self.frame == self.frames: # end of video
|
||||
self.count += 1
|
||||
self.cap.release()
|
||||
else:
|
||||
# Move to the next file if the current video ended or failed to open
|
||||
self.count += 1
|
||||
if self.cap:
|
||||
self.cap.release()
|
||||
if self.count < self.nf:
|
||||
self._new_video(self.files[self.count])
|
||||
else:
|
||||
self.mode = "image"
|
||||
im0 = cv2.imread(path) # BGR
|
||||
if im0 is None:
|
||||
raise FileNotFoundError(f"Image Not Found {path}")
|
||||
paths.append(path)
|
||||
imgs.append(im0)
|
||||
info.append(f"image {self.count + 1}/{self.nf} {path}: ")
|
||||
self.count += 1 # move to the next file
|
||||
if self.count >= self.ni: # end of image list
|
||||
break
|
||||
|
||||
return paths, imgs, info
|
||||
|
||||
def _new_video(self, path):
|
||||
"""Create a new video capture object."""
|
||||
"""Creates a new video capture object for the given path."""
|
||||
self.frame = 0
|
||||
self.cap = cv2.VideoCapture(path)
|
||||
self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
|
||||
if not self.cap.isOpened():
|
||||
raise FileNotFoundError(f"Failed to open video {path}")
|
||||
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
|
||||
|
||||
def __len__(self):
|
||||
"""Returns the number of files in the object."""
|
||||
return self.nf # number of files
|
||||
"""Returns the number of batches in the object."""
|
||||
return math.ceil(self.nf / self.bs) # number of files
|
||||
|
||||
|
||||
class LoadPilAndNumpy:
|
||||
"""
|
||||
Load images from PIL and Numpy arrays for batch processing.
|
||||
|
||||
def __init__(self, im0, imgsz=640):
|
||||
This class is designed to manage loading and pre-processing of image data from both PIL and Numpy formats.
|
||||
It performs basic validation and format conversion to ensure that the images are in the required format for
|
||||
downstream processing.
|
||||
|
||||
Attributes:
|
||||
paths (list): List of image paths or autogenerated filenames.
|
||||
im0 (list): List of images stored as Numpy arrays.
|
||||
mode (str): Type of data being processed, defaults to 'image'.
|
||||
bs (int): Batch size, equivalent to the length of `im0`.
|
||||
|
||||
Methods:
|
||||
_single_check(im): Validate and format a single image to a Numpy array.
|
||||
"""
|
||||
|
||||
def __init__(self, im0):
|
||||
"""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'
|
||||
# Generate fake paths
|
||||
self.mode = "image"
|
||||
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
|
||||
@ -310,7 +430,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, [""] * self.bs
|
||||
|
||||
def __iter__(self):
|
||||
"""Enables iteration for class LoadPilAndNumpy."""
|
||||
@ -319,18 +439,36 @@ class LoadPilAndNumpy:
|
||||
|
||||
|
||||
class LoadTensor:
|
||||
"""
|
||||
Load images from torch.Tensor data.
|
||||
|
||||
This class manages the loading and pre-processing of image data from PyTorch tensors for further processing.
|
||||
|
||||
Attributes:
|
||||
im0 (torch.Tensor): The input tensor containing the image(s).
|
||||
bs (int): Batch size, inferred from the shape of `im0`.
|
||||
mode (str): Current mode, set to 'image'.
|
||||
paths (list): List of image paths or filenames.
|
||||
count (int): Counter for iteration, initialized at 0 during `__iter__()`.
|
||||
|
||||
Methods:
|
||||
_single_check(im, stride): Validate and possibly modify the input tensor.
|
||||
"""
|
||||
|
||||
def __init__(self, im0) -> None:
|
||||
"""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)
|
||||
@ -338,9 +476,11 @@ class LoadTensor:
|
||||
im = im.unsqueeze(0)
|
||||
if im.shape[2] % stride or im.shape[3] % stride:
|
||||
raise ValueError(s)
|
||||
if im.max() > 1.0:
|
||||
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.')
|
||||
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."
|
||||
)
|
||||
im = im.float() / 255.0
|
||||
|
||||
return im
|
||||
@ -355,7 +495,7 @@ class LoadTensor:
|
||||
if self.count == 1:
|
||||
raise StopIteration
|
||||
self.count += 1
|
||||
return self.paths, self.im0, None, ''
|
||||
return self.paths, self.im0, [""] * self.bs
|
||||
|
||||
def __len__(self):
|
||||
"""Returns the batch size."""
|
||||
@ -363,26 +503,23 @@ class LoadTensor:
|
||||
|
||||
|
||||
def autocast_list(source):
|
||||
"""
|
||||
Merges a list of source of different types into a list of numpy arrays or PIL images
|
||||
"""
|
||||
"""Merges a list of source of different types into a list of numpy arrays or PIL images."""
|
||||
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
|
||||
|
||||
|
||||
LOADERS = LoadStreams, LoadPilAndNumpy, LoadImages, LoadScreenshots # tuple
|
||||
|
||||
|
||||
def get_best_youtube_url(url, use_pafy=False):
|
||||
def get_best_youtube_url(url, use_pafy=True):
|
||||
"""
|
||||
Retrieves the URL of the best quality MP4 video stream from a given YouTube video.
|
||||
|
||||
@ -397,16 +534,22 @@ def get_best_youtube_url(url, use_pafy=False):
|
||||
(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")
|
||||
|
||||
|
||||
# Define constants
|
||||
LOADERS = (LoadStreams, LoadPilAndNumpy, LoadImagesAndVideos, LoadScreenshots)
|
||||
|
Reference in New Issue
Block a user