| 366 | |
| 367 | |
| 368 | def preprocess(dataset_path, mask_path, mode, num_frames, stride, logger): |
| 369 | # Define paths to videos in dataset |
| 370 | movies_path_list = sorted([Path(p) for p in glob.glob(os.path.join(dataset_path, '**/*.mp4'), recursive=True)]) |
| 371 | if len(movies_path_list) == 0: |
| 372 | logger.error(f"No videos found in {dataset_path}") |
| 373 | sys.exit() |
| 374 | logger.info(f"{len(movies_path_list)} videos found in {dataset_path}") |
| 375 | |
| 376 | # Define paths to masks in dataset |
| 377 | if mask_path is not None: |
| 378 | masks_path_list = sorted([Path(p) for p in glob.glob(os.path.join(mask_path, '**/*.mp4'), recursive=True)]) |
| 379 | if len(masks_path_list) == 0: |
| 380 | logger.error(f"No masks found in {mask_path}") |
| 381 | # sys.exit() |
| 382 | logger.info(f"{len(masks_path_list)} masks found in {mask_path}") |
| 383 | |
| 384 | # Start timer |
| 385 | start_time = time.monotonic() |
| 386 | |
| 387 | # Define the number of processes based on CPU capabilities |
| 388 | num_processes = os.cpu_count() |
| 389 | |
| 390 | # Use multiprocessing to process videos in parallel |
| 391 | with concurrent.futures.ThreadPoolExecutor(max_workers=num_processes) as executor: |
| 392 | futures = [] |
| 393 | for movie_path in movies_path_list: |
| 394 | # Check if there is a mask for the video |
| 395 | if mask_path is not None: |
| 396 | if movie_path.stem not in [path.stem for path in masks_path_list]: |
| 397 | logger.error(f"No mask for video {movie_path}") |
| 398 | # Define the mask path |
| 399 | mask_path = next((path for path in masks_path_list if path.stem == movie_path.stem), None) |
| 400 | if mask_path is None: |
| 401 | logger.error(f"Mask path not found for video {movie_path}") |
| 402 | # Create a future for each video and submit it for processing |
| 403 | futures.append( |
| 404 | executor.submit( |
| 405 | video_manipulate, |
| 406 | movie_path, |
| 407 | mask_path, |
| 408 | dataset_path, |
| 409 | mode, |
| 410 | num_frames, |
| 411 | stride, |
| 412 | ) |
| 413 | ) |
| 414 | # Wait for all futures to complete and log any errors |
| 415 | for future in tqdm(concurrent.futures.as_completed(futures), total=len(movies_path_list)): |
| 416 | # Print the current time |
| 417 | logger.info(f"Current time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") |
| 418 | try: |
| 419 | future.result() |
| 420 | except Exception as e: |
| 421 | logger.error(f"Error processing video: {e}") |
| 422 | |
| 423 | # End timer |
| 424 | end_time = time.monotonic() |
| 425 | duration_minutes = (end_time - start_time) / 60 |