| 246 | |
| 247 | |
| 248 | class StreamDataLoader: |
| 249 | def __init__(self, data_cfg: DataConfig): |
| 250 | self.source = data_cfg.source |
| 251 | self.running = True |
| 252 | self.is_stream = isinstance(self.source, int) or str(self.source).lower().startswith("rtmp://") |
| 253 | |
| 254 | self.transform = AugmentationComposer([], data_cfg.image_size) |
| 255 | self.stop_event = Event() |
| 256 | |
| 257 | if self.is_stream: |
| 258 | import cv2 |
| 259 | |
| 260 | self.cap = cv2.VideoCapture(self.source) |
| 261 | else: |
| 262 | self.source = Path(self.source) |
| 263 | self.queue = Queue() |
| 264 | self.thread = Thread(target=self.load_source) |
| 265 | self.thread.start() |
| 266 | |
| 267 | def load_source(self): |
| 268 | if self.source.is_dir(): # image folder |
| 269 | self.load_image_folder(self.source) |
| 270 | elif any(self.source.suffix.lower().endswith(ext) for ext in [".mp4", ".avi", ".mkv"]): # Video file |
| 271 | self.load_video_file(self.source) |
| 272 | else: # Single image |
| 273 | self.process_image(self.source) |
| 274 | |
| 275 | def load_image_folder(self, folder): |
| 276 | folder_path = Path(folder) |
| 277 | for file_path in folder_path.rglob("*"): |
| 278 | if self.stop_event.is_set(): |
| 279 | break |
| 280 | if file_path.suffix.lower() in [".jpg", ".jpeg", ".png", ".bmp"]: |
| 281 | self.process_image(file_path) |
| 282 | |
| 283 | def process_image(self, image_path): |
| 284 | image = Image.open(image_path).convert("RGB") |
| 285 | if image is None: |
| 286 | raise ValueError(f"Error loading image: {image_path}") |
| 287 | self.process_frame(image) |
| 288 | |
| 289 | def load_video_file(self, video_path): |
| 290 | import cv2 |
| 291 | |
| 292 | cap = cv2.VideoCapture(str(video_path)) |
| 293 | while self.running: |
| 294 | ret, frame = cap.read() |
| 295 | if not ret: |
| 296 | break |
| 297 | self.process_frame(frame) |
| 298 | cap.release() |
| 299 | |
| 300 | def process_frame(self, frame): |
| 301 | if isinstance(frame, np.ndarray): |
| 302 | # TODO: we don't need cv2 |
| 303 | import cv2 |
| 304 | |
| 305 | frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
no outgoing calls