Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| import shutil | |
| import tempfile | |
| import zipfile | |
| import gradio as gr | |
| def process_zips(camera_zips, regex_pattern, start_frame, end_frame): | |
| """ | |
| Processes uploaded ZIP files containing frame images from one or more cameras. | |
| Args: | |
| camera_zips (list): A list of file paths for the uploaded ZIP files (each representing one camera folder). | |
| regex_pattern (str): A regex pattern with one capturing group that extracts the numeric frame value. | |
| start_frame (str): The start frame number (as string, will be converted to int). | |
| end_frame (str): The end frame number (as string, will be converted to int). | |
| Returns: | |
| str: The path to the output ZIP file containing the sorted frames. | |
| The function: | |
| 1. Extracts each uploaded ZIP into a temporary folder. | |
| 2. Uses the ZIP file name (without extension) as the camera name. | |
| 3. Iterates over every file in the extracted folder, applying the regex to extract the numeric frame. | |
| 4. For files whose numeric value is between start_frame and end_frame (inclusive), | |
| computes the frame index as: | |
| frame_index = extracted_number - start_frame + 1 | |
| 5. Copies each matching file into an output folder structure where each frame index gets its own folder. | |
| 6. Compresses the entire output directory into a ZIP file and returns its path. | |
| """ | |
| try: | |
| start_frame = int(start_frame) | |
| end_frame = int(end_frame) | |
| except ValueError: | |
| return "Error: Start and End frame numbers must be integers." | |
| if start_frame > end_frame: | |
| return "Error: Start frame must be less than or equal to End frame." | |
| total_frames = end_frame - start_frame + 1 | |
| # Create a temporary output directory for sorted frames | |
| output_dir = tempfile.mkdtemp(prefix="sorted_frames_") | |
| # List to hold (camera_name, folder) tuples | |
| camera_folders = [] | |
| for file_path in camera_zips: | |
| zip_name = os.path.basename(file_path) | |
| # Create a temporary directory to extract the ZIP contents | |
| extract_dir = os.path.join(tempfile.gettempdir(), zip_name + "_extracted") | |
| os.makedirs(extract_dir, exist_ok=True) | |
| with zipfile.ZipFile(file_path, "r") as zip_ref: | |
| zip_ref.extractall(extract_dir) | |
| # If the ZIP extracts to a single subdirectory, use that folder | |
| subdirs = [os.path.join(extract_dir, d) for d in os.listdir(extract_dir) | |
| if os.path.isdir(os.path.join(extract_dir, d))] | |
| camera_folder = subdirs[0] if len(subdirs) == 1 else extract_dir | |
| # Use the ZIP filename (without extension) as the camera name | |
| camera_name = os.path.splitext(zip_name)[0] | |
| camera_folders.append((camera_name, camera_folder)) | |
| # Build a mapping for each camera: { frame_index: file_full_path } | |
| camera_mappings = {} | |
| for camera_name, folder in camera_folders: | |
| mapping = {} | |
| for file_name in os.listdir(folder): | |
| full_path = os.path.join(folder, file_name) | |
| if os.path.isfile(full_path): | |
| try: | |
| m = re.search(regex_pattern, file_name) | |
| except re.error as e: | |
| print(f"Regex error on file {file_name}: {e}") | |
| continue | |
| if m: | |
| try: | |
| num = int(m.group(1)) | |
| except ValueError: | |
| continue | |
| if start_frame <= num <= end_frame: | |
| frame_index = num - start_frame + 1 | |
| mapping[frame_index] = full_path | |
| camera_mappings[camera_name] = mapping | |
| # Create output subfolders for each frame index and copy matching files from each camera | |
| for frame_index in range(1, total_frames + 1): | |
| frame_folder = os.path.join(output_dir, f"Frame_{frame_index}") | |
| os.makedirs(frame_folder, exist_ok=True) | |
| for camera_name, mapping in camera_mappings.items(): | |
| if frame_index in mapping: | |
| src_file = mapping[frame_index] | |
| file_base = os.path.basename(src_file) | |
| # Rename the file by prefixing it with the camera name | |
| dst_file = os.path.join(frame_folder, f"{camera_name}_{file_base}") | |
| try: | |
| shutil.copy(src_file, dst_file) | |
| except Exception as e: | |
| print(f"Error copying {src_file} to {dst_file}: {e}") | |
| # Compress the output directory into a ZIP file | |
| output_zip_base = os.path.join(tempfile.gettempdir(), "sorted_frames") | |
| shutil.make_archive(base_name=output_zip_base, format="zip", root_dir=output_dir) | |
| output_zip_path = output_zip_base + ".zip" | |
| return output_zip_path | |
| # Create the Gradio interface using type "filepath" | |
| iface = gr.Interface( | |
| fn=process_zips, | |
| inputs=[ | |
| gr.File(label="Camera Folder ZIPs", file_count="multiple", type="filepath"), | |
| gr.Textbox(label="Regex Pattern (with one capturing group)"), | |
| gr.Textbox(label="Start Frame Number"), | |
| gr.Textbox(label="End Frame Number") | |
| ], | |
| outputs=gr.File(label="Download Sorted Frames ZIP"), | |
| title="Frame Sorting Tool", | |
| description=( | |
| "Upload one or more ZIP files (each representing a camera folder with frame files), " | |
| "enter a regex pattern (for example, `_(\\d+)`), and provide the start and end frame numbers. " | |
| "The tool will group frames by frame index (so that frame 1 from all cameras is together) " | |
| "and return a downloadable ZIP file with the sorted frames." | |
| "Mention the Readme for further information on Regex pattern and Start and End Frames format." | |
| ) | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |