Spaces:
Configuration error
Configuration error
File size: 2,365 Bytes
dcd1471 |
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 |
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")
|