Spaces:
Runtime error
Runtime error
| import os | |
| import cv2 | |
| import numpy as np | |
| import argparse | |
| from PIL import Image | |
| def convert_and_copy_masks(src_folder, dest_folder, index_offset, start_frame, end_frame): | |
| # Ensure that the destination folder exists | |
| if not os.path.exists(dest_folder): | |
| os.makedirs(dest_folder) | |
| # Iterate over all files in the source folder | |
| for filename in os.listdir(src_folder): | |
| if filename.endswith(".png"): # Process only PNG files | |
| # Extract the numeric part of the file name | |
| file_index = int(filename.split('.')[0]) | |
| # Check if the file index is within the specified frame range | |
| if start_frame <= file_index <= end_frame: | |
| # Calculate the new index | |
| new_index = file_index + index_offset | |
| # Construct the new file name | |
| new_filename = f"{new_index:05d}.png" | |
| # Read the original mask image (grayscale) | |
| src_path = os.path.join(src_folder, filename) | |
| mask = cv2.imread(src_path, cv2.IMREAD_GRAYSCALE) | |
| # Apply thresholding to convert the mask to 0 and 1 | |
| _, mask = cv2.threshold(mask, 0.5, 1, cv2.THRESH_BINARY) | |
| # Convert mask value range to 0-255 | |
| mask = (mask * 255).astype(np.uint8) | |
| # Save the converted mask to the destination folder | |
| dest_path = os.path.join(dest_folder, new_filename) | |
| cv2.imwrite(dest_path, mask) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Convert and copy mask images with index adjustment") | |
| parser.add_argument('--src_folder', type=str, required=True, help="Path to the source folder containing mask images") | |
| parser.add_argument('--dest_folder', type=str, required=True, help="Path to the destination folder for converted masks") | |
| parser.add_argument('--index_offset', type=int, default=0, help="Index offset to apply to file names") | |
| parser.add_argument('--start_frame', type=int, default=0, help="Start frame number") | |
| parser.add_argument('--end_frame', type=int, required=True, help="End frame number") | |
| args = parser.parse_args() | |
| convert_and_copy_masks(args.src_folder, args.dest_folder, args.index_offset, args.start_frame, args.end_frame) | |
| print("Mask conversion and copying completed") | |