Spaces:
Build error
Build error
File size: 4,609 Bytes
e1832f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import os
import cv2
import glob
import math
import numpy as np
from pathlib import Path
from PIL import Image
VID_FORMATS = "asf", "avi", "gif", "m4v", "mkv", "mov", "mp4", "mpeg", "mpg", "ts", "wmv" # include video suffixes
class LoadImagesAndVideos:
"""
A data loader for handling both images and videos, providing batches of frames or images for processing.
Supports various image formats, including HEIC, and handles text files with paths to images/videos.
"""
def __init__(self, path, batch_size=1, vid_stride=1):
self.batch_size = batch_size
self.vid_stride = vid_stride
self.files = self._load_files(path)
self.video_flag = [self._is_video(f) for f in self.files]
self.nf = len(self.files)
self.ni = sum(not is_video for is_video in self.video_flag)
self.mode = "image"
self.cap = None
if any(self.video_flag):
self._start_video(self.files[self.video_flag.index(True)])
if not self.files:
raise FileNotFoundError(f"No images or videos found in {path}.")
def _load_files(self, path):
"""Load files from a given path, which may be a directory, list, or text file."""
if isinstance(path, str) and Path(path).suffix == ".txt":
path = Path(path).read_text().splitlines()
files = []
for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
p = str(Path(p).absolute())
if "*" in p:
files.extend(glob.glob(p, recursive=True))
elif os.path.isdir(p):
files.extend(glob.glob(os.path.join(p, "*.*")))
elif os.path.isfile(p):
files.append(p)
else:
raise FileNotFoundError(f"{p} does not exist")
return files
def _is_video(self, file_path):
"""Check if a file is a video based on its extension."""
return file_path.split('.')[-1].lower() in VID_FORMATS
def __iter__(self):
self.count = 0
return self
def __next__(self):
paths, imgs, infos = [], [], []
while len(imgs) < self.batch_size:
if self.count >= self.nf:
if imgs:
return paths, imgs, infos
else:
raise StopIteration
path = self.files[self.count]
if self.video_flag[self.count]:
self._process_video(paths, imgs, infos, path)
else:
self._process_image(paths, imgs, infos, path)
self.count += 1
return paths, imgs, infos
def _process_image(self, paths, imgs, infos, path):
"""Process an image file and append it to the batch."""
img = self._read_image(path)
if img is not None:
paths.append(path)
imgs.append(img)
infos.append(f"image {self.count + 1}/{self.nf} {path}")
def _process_video(self, paths, imgs, infos, path):
"""Process a video file, reading frames as per the stride."""
self.mode = "video"
if not self.cap or not self.cap.isOpened():
self._start_video(path)
success = False
for _ in range(self.vid_stride):
success = self.cap.grab()
if not success:
break
if success:
_, frame = self.cap.retrieve()
paths.append(path)
imgs.append(frame)
infos.append(f"video {self.count + 1}/{self.nf} frame {self.frame}/{self.frames} {path}")
self.frame += 1
if self.frame >= self.frames:
self.cap.release()
def _read_image(self, path):
"""Read an image from a file, handling HEIC format if necessary."""
if path.lower().endswith("heic"):
from pillow_heif import register_heif_opener
register_heif_opener()
with Image.open(path) as img:
return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
else:
return cv2.imread(path)
def _start_video(self, path):
"""Initialize video capture for a new video file."""
self.cap = cv2.VideoCapture(path)
if not self.cap.isOpened():
raise FileNotFoundError(f"Failed to open video {path}")
self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
self.frame = 0
def __len__(self):
return math.ceil(self.nf / self.batch_size)
|