HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.1.26 LoadImagesAndVideos batched inference (#8817)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-03-10 19:15:44 +01:00 • committed by GitHub
parent 1f9667fff2
commit 7451ca1f54
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 186 additions and 171 deletions
Download patch file Download diff file Expand all files Collapse all files

125
ultralytics/data/loaders.py
View file

@ -24,7 +24,7 @@ from ultralytics.utils.checks import check_requirements
class SourceTypes:
"""Class to represent various types of input sources for predictions."""
webcam: bool = False
stream: bool = False
screenshot: bool = False
from_img: bool = False
tensor: bool = False
@ -32,9 +32,7 @@ class SourceTypes:
class LoadStreams:
"""
Stream Loader for various types of video streams.
Suitable for use with `yolo predict source='rtsp://example.com/media.mp4'`, supports RTSP, RTMP, HTTP, and TCP streams.
Stream Loader for various types of video streams, Supports RTSP, RTMP, HTTP, and TCP streams.
Attributes:
sources (str): The source input paths or URLs for the video streams.
@ -57,6 +55,11 @@ class LoadStreams:
__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):
@ -69,6 +72,7 @@ class LoadStreams:
sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
n = len(sources)
self.bs = n
self.fps = [0] * n # frames per second
self.frames = [0] * n
self.threads = [None] * n
@ -76,6 +80,8 @@ class LoadStreams:
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
self.info = [""] * n
self.is_video = [True] * n
for i, s in enumerate(sources): # index, source
# Start thread to read frames from video stream
st = f"{i + 1}/{n}: {s}... "
@ -109,9 +115,6 @@ class LoadStreams:
self.threads[i].start()
LOGGER.info("") # newline
# Check for common shapes
self.bs = self.__len__()
def update(self, i, cap, stream):
"""Read stream `i` frames in daemon thread."""
n, f = 0, self.frames[i] # frame number, frame array
@ -175,11 +178,11 @@ 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, self.is_video, self.info
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:
@ -243,10 +246,10 @@ class LoadScreenshots:
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], [True], [s] # screen, img, is_video, string
class LoadImages:
class LoadImagesAndVideos:
"""
YOLOv8 image/video dataloader.
@ -269,7 +272,7 @@ class LoadImages:
_new_video(path): Create a new cv2.VideoCapture object for a given video path.
"""
def __init__(self, path, vid_stride=1):
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
@ -298,7 +301,7 @@ class LoadImages:
self.video_flag = [False] * ni + [True] * nv
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:
@ -315,49 +318,68 @@ 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, is_video, info = [], [], [], []
while len(imgs) < self.bs:
if self.count >= self.nf: # end of file list
if len(imgs) > 0:
return paths, imgs, is_video, 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)
is_video.append(True)
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)
is_video.append(False) # no capture object for images
info.append(f"image {self.count + 1}/{self.nf} {path}: ")
self.count += 1 # move to the next file
return paths, imgs, is_video, 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)
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:
@ -373,7 +395,6 @@ class LoadPilAndNumpy:
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`.
count (int): Counter for iteration, initialized at 0 during `__iter__()`.
Methods:
_single_check(im): Validate and format a single image to a Numpy array.
@ -386,7 +407,6 @@ class LoadPilAndNumpy:
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.mode = "image"
# Generate fake paths
self.bs = len(self.im0)
@staticmethod
@ -409,7 +429,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, [False] * self.bs, [""] * self.bs
def __iter__(self):
"""Enables iteration for class LoadPilAndNumpy."""
@ -474,7 +494,7 @@ class LoadTensor:
if self.count == 1:
raise StopIteration
self.count += 1
return self.paths, self.im0, None, ""
return self.paths, self.im0, [False] * self.bs, [""] * self.bs
def __len__(self):
"""Returns the batch size."""
@ -498,9 +518,6 @@ def autocast_list(source):
return files
LOADERS = LoadStreams, LoadPilAndNumpy, LoadImages, LoadScreenshots # tuple
def get_best_youtube_url(url, use_pafy=True):
"""
Retrieves the URL of the best quality MP4 video stream from a given YouTube video.
@ -531,3 +548,7 @@ def get_best_youtube_url(url, use_pafy=True):
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)