# Project export: SnapAR3D This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: TreeHacks 2025 - Tagline: SnapAR3D uses AR to guide optimal photo capture and record precise camera poses relative to your object. We convert this data for nvdifrec to create accurate 3D models. - Devpost: https://devpost.com/software/snapar3d - GitHub: https://github.com/wwangg22/treehcks - Video: https://www.youtube.com/embed/9cUk6gcapfg?enablejsapi=1&hl=en_US&rel=0&start=&version=3&wmode=transparent - Team: 1 GitHub contributor(s) — william (1 commits) ## Devpost submission (written by the team) ### Inspiration We were inspired by the increasing demand for accessible 3D modeling tools and the incredible potential of augmented reality. Seeing how AR could guide users to capture high-quality photos, we envisioned a tool that bridges real-world imagery and accurate 3D reconstructions. ### What it does SnapAR3D uses AR to guide you in capturing optimal photos of an object, while precisely recording the camera’s positions and orientations relative to that object. This data is transformed into matrices and imported into NVIDIA Diffractive Reconstruction (nvdifrec) to generate high-quality 3D models. ### How we built it We built SnapAR3D by integrating ARKit to capture photos and sensor data, and COLMAP to format the camera parameters. A custom conversion pipeline transforms the captured data into NeRF-compatible transformation matrices, which nvdifrec uses for reconstruction. We also developed a lightweight web host for easy sharing of the resulting transforms.json file. ### Challenges we ran into One of our biggest challenges was fine tuning—adjusting the photo capture process and data transformation parameters to ensure that the final 3D model was as accurate and detailed as possible. We also had to address coordinate system discrepancies between ARKit and NeRF, and seamlessly integrate data from multiple sources. ### Accomplishments we're proud of Successfully guiding users with AR to capture optimal images. Training a model that is comparable in quality and much faster in training time ### What we learned We deepened our understanding in AR app development, gained valuable insights into how 3D models are built using 2D photos, and learned to always push through the challenges because it will be worth it in the end. ### What's next Next, we plan to improve the user interface for an even smoother capture experience, expand support for diverse objects and scenes, and fine tune the model further while enhancing real-time feedback during data capture. ## README (from the GitHub repository) No README available. ## Detected evidence (automated analysis) Indexed codebase: 10 recognized source files, 29 KB. - Python (language) — detected in the code - JavaScript (language) — claimed on Devpost, not found in the code - Swift (language) — claimed on Devpost, not found in the code ## Codebase structure (from repository index) ### Files (12 of 12) ``` .gitignore .gitmodules advanced_server.py basic_server.py convert_my_data_to_transforms.py data_processing_utils/cleanup.py data_processing_utils/fixname.py data_processing_utils/removepng.py data_processing_utils/scale.py data_processing_utils/split_for_test.py data_processing_utils/split_for_val.py data_processing.py ``` ### Dependencies No dependency index available. ### Recent commits (newest first) - first ## Key source files (fetched from GitHub, selected and truncated for size) ### basic_server.py ```python import os import re import time from http.server import BaseHTTPRequestHandler, HTTPServer class SimpleFileUploadHandler(BaseHTTPRequestHandler): def do_POST(self): # 1. Get the content length from headers content_length = int(self.headers['Content-Length']) # 2. Read the file data from the request body file_data = self.rfile.read(content_length) # 3. Try to extract the original filename from the Content-Disposition header filename = f"uploaded_{int(time.time())}" content_disp = self.headers.get('Content-Disposition') if content_disp: # Look for a filename="..." match = re.search(r'filename="([^"]+)"', content_disp) if match: original_filename = match.group(1) # Extract extension (if any) from the original filename _, ext = os.path.splitext(original_filename) filename += ext # 4. Save the file data to disk with open(filename, 'wb') as f: f.write(file_data) # 5. Send a response back to the client self.send_response(200) self.end_headers() self.wfile.write(b"File uploaded successfully.") def run_server(): server_address = ('', 8000) # Listen on all interfaces, port 8000 httpd = HTTPServer(server_address, SimpleFileUploadHandler) print("Serving on port 8000...") httpd.serve_forever() if __name__ == '__main__': run_server() ``` ### convert_my_data_to_transforms.py ```python #!/usr/bin/env python3 import os import json import base64 import numpy as np import math import cv2 def qvec2rotmat(qvec): """ Convert a quaternion [q0, q1, q2, q3] (with q0 as the scalar) to a rotation matrix. """ q0, q1, q2, q3 = qvec return np.array([ [1 - 2*(q2**2) - 2*(q3**2), 2*q1*q2 - 2*q3*q0, 2*q1*q3 + 2*q2*q0], [2*q1*q2 + 2*q3*q0, 1 - 2*(q1**2) - 2*(q3**2), 2*q2*q3 - 2*q1*q0], [2*q1*q3 - 2*q2*q0, 2*q2*q3 + 2*q1*q0, 1 - 2*(q1**2) - 2*(q2**2)] ]) def compute_transform(qvec, tvec): """ Compute the camera-to-world transformation matrix. Uses a COLMAP-to-NeRF conversion approach: 1. Negates the quaternion. 2. Forms the 4x4 matrix m = [R, t; 0,0,0,1] with R as rotation. 3. Inverts m to obtain the camera-to-world matrix. """ qvec_neg = -np.array(qvec) R = qvec2rotmat(qvec_neg) t = np.array(tvec).reshape((3, 1)) bottom = np.array([[0, 0, 0, 1]]) m = np.concatenate([np.concatenate([R, t], axis=1), bottom], axis=0) c2w = np.linalg.inv(m) return c2w def sharpness(imagePath): """ Compute the sharpness of an image using the variance of its Laplacian. Rotates the image 90° clockwise before processing. """ image = cv2.imread(imagePath) image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE) if image is None: return 0.0 gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) return cv2.Laplacian(gray, cv2.CV_64F).var() def convert_my_data_to_transforms(json_path, out_dir="tmp_data", keep_colmap_coords=False): """ Convert the input JSON data (expected as a list of records) into a NeRF-formatted transforms JSON file. Parameters: json_path (str): Path to the JSON file containing the data. out_dir (str): Directory where output images and the transforms JSON will be saved. keep_colmap_coords (bool): If False, reorient the transforms as per the NVIDIA script. Raises: FileNotFoundError: If json_path does not exist. ValueError: If the JSON data is not a non-empty list. """ if not os.path.isfile(json_path): raise FileNotFoundError(f"Error: {json_path} does not exist.") with open(json_path, 'r') as f: data = json.load(f) if not isinstance(data, list) or len(data) == 0: raise ValueError("JSON must be a non-empty list.") os.makedirs(out_dir, exist_ok=True) images_dir = os.path.join(out_dir, "images") os.makedirs(images_dir, exist_ok=True) # Use the intrinsics from the first record. first = data[0] intr = first.get("cameraIntrinsics", [1500, 1500, 960, 540]) # [fx, fy, cx, cy] fx, fy, cx, cy = intr # Guess resolution as width = 2*cx, height = 2*cy. width = int(cx * 2) height = int(cy * 2) camera_angle_x = 2 * math.atan(width / (2 * fx)) # Prepare output dictionary for transforms in NeRF format. out = { "camera_angle_x": camera_angle_x, "fl_x": fx, "fl_y": fy, "cx": cx, "cy": cy, "w": width, "h": height, "aabb_scale": 16, # Adjust as needed "frames": [] } # Process each record in the JSON. for i, rec in enumerate(data, start=1): # Save the image from imageBase64 into images_dir. b64_str = rec.get("imageBase64", "") img_name = f"patch_image_{i:04d}.png" img_path = os.path.join(images_dir, img_name) if b64_str: png_data = base64.b64decode(b64_str) with open(img_path, 'wb') as img_file: img_file.write(png_data) else: print(f"Warning: No imageBase64 for record {i}") # Compute sharpness. s = sharpness(img_path) # Get the transform from cameraOrientation and cameraPosition. qvec = rec["cameraOrientation"] # [qw, qx, qy, qz] tvec = rec["cameraPosition"] # [tx, ty, tz] c2w = compute_transform(qvec, tvec) # Optionally, reorient the transforms if not keeping COLMAP coordinates. if not keep_colmap_coords: # Flip the z axis (as in the NVIDIA script) c2w[0:3,2] *= -1 # c2w[0:3,1] *= -1 # Uncomment if needed. c2w = c2w[[1, 0, 2, 3], :] c2w[2, :] *= -1 frame = { "file_path": "./train/" + img_name.split('.')[0], "sharpness": s, "transform_matrix": c2w.tolist() } out["frames"].append(frame) transforms_path = os.path.join(out_dir, "transforms_train.json") with open(transforms_path, "w") as outfile: json.dump(out, outfile, indent=2) print(f"Successfully wrote transforms.json to {transforms_path}") # If you want to run this module directly. if __name__ == "__main__": import sys if len(sys.argv) < 2: print("Usage: python convert_my_data_to_transforms.py [output_dir] [--keep_colmap_coords]") sys.exit(1) json_path = sys.argv[1] out_dir = sys.argv[2] if len(sys.argv) > 2 else "tmp_data" keep_coords = "--keep_colmap_coords" in sys.argv convert_my_data_to_transforms(json_path, out_dir, keep_coords) ``` ### advanced_server.py ```python import os import re import time import msvcrt import json import threading import queue from http.server import BaseHTTPRequestHandler, HTTPServer from rembg import remove import uuid from data_processing_utils.scale import scale_images from data_processing_utils.split_for_test import split_frames_into_test_set from data_processing_utils.split_for_val import split_frames_into_val_set from convert_my_data_to_transforms import convert_my_data_to_transforms from nvdiffrec_test.train_export import run_nvdiffrec # Global queue to store JSON POST requests. request_queue = queue.Queue() class AdvancedRequestHandler(BaseHTTPRequestHandler): def do_POST(self): """Handle POST requests: verify JSON, enqueue it, and return a response.""" try: content_length = int(self.headers.get('Content-Length', 0)) except (ValueError, TypeError): self.send_response(411) self.end_headers() self.wfile.write(b"Missing or invalid Content-Length header") return # Read the data sent in the request. post_data = self.rfile.read(content_length) # Ensure the Content-Type header indicates JSON. content_type = self.headers.get('Content-Type', '') if 'application/json' not in content_type.lower(): self.send_response(400) self.end_headers() self.wfile.write(b"Only application/json content is supported.") return try: data = json.loads(post_data.decode('utf-8')) except json.JSONDecodeError: self.send_response(400) self.end_headers() self.wfile.write(b"Invalid JSON provided.") return # Write JSON data to a file in a local directory. local_dir = "json_data" # Directory where JSON files will be stored. os.makedirs(local_dir, exist_ok=True) # Create a unique filename for the JSON file. filename = f"data_{uuid.uuid4().hex}.json" file_path = os.path.join(local_dir, filename) with open(file_path, "w") as f: json.dump(data, f, indent=2) # Enqueue the path to the JSON file. request_queue.put(file_path) # Send a success response. self.send_response(200) self.end_headers() self.wfile.write(b"JSON received and queued.") def run_server(): """Start the HTTP server.""" server_address = ('', 8000) # Listen on all interfaces, port 8000. httpd = HTTPServer(server_address, AdvancedRequestHandler) print("Serving on port 8000...") httpd.serve_forever() def process_video_request(): """ Simulate processing a video request. In this example, we dequeue all queued JSON requests and print them. """ request_queue.queue(9) def get_input_non_blocking(): """Check if a key has been pressed on Windows.""" if msvcrt.kbhit(): # Read a single character (as a Unicode string) return msvcrt.getwch().lower() return None def main(): """Continuously listen for user input from the command prompt.""" while True: command = get_input_non_blocking() if command: # Since msvcrt.getwch() returns one character at a time, # you might need to implement a buffer if multi-character commands are required. if command == 'v': process_video_request() else: if not request_queue.empty(): front_item = request_queue.get() print("Removed front item:", front_item) else: print("The queue is empty.") #call processing file to process JSON file #this function saves the images in nerf format as folder/images and folder/transform.json #we need to first remove the background from the images, resize it, rename the images to train, #split 5 photos to test, setup config so it reflects the width, height convert_my_data_to_transforms(front_item, "tmp_data") #this is for removing the background, editing IN PLACE folder = './tmp_data/train' folder_noimg = './tmp/data' for filename in os.listdir(folder): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): path = os.path.join(folder, filename) with open(path, 'rb') as f: output = remove(f.read()) with open(path, 'wb') as f: f.write(output) print(f"Processed: {filename}") dims = scale_images(folder, folder, 0.5, 0.5) W, H = dims[0] split_frames_into_test_set(folder_noimg+'/transforms_train.json',folder_noimg+'/transforms_test.json', test_folder=folder_noimg + '/test', num_test_frames=5) config = { "ref_mesh": "./tmp_data", "random_textures": True, "iter": 4000, "save_interval": 100, "texture_res": [ 2048, 2048 ], "train_res": [H, W], "batch": 6, "learning_rate": [0.03, 0.01], "ks_min" : [0, 0.25, 0.0], "dmtet_grid" : 32, "mesh_scale" : 2.4, "laplace_scale" : 9000, "display": [{"latlong" : True}, {"bsdf" : "kd"}, {"bsdf" : "ks"}, {"bsdf" : "normal"}], "background" : "white", "validate": False, "out_dir": "./tmp_out" } run_nvdiffrec(config) if __name__ == '__main__': # Start the HTTP server in a separate daemon thread. server_thread = threading.Thread(target=run_se [truncated — 115 more characters] ``` ### data_processing_utils/fixname.py ```python #!/usr/bin/env python3 import json import sys def replace_json_frame_paths( input_json_path: str, output_json_path: str, old_string: str = "./images/", new_string: str = "./train/" ): """ Reads JSON data from input_json_path and replaces occurrences of old_string with new_string in the 'file_path' field of each frame, then writes the modified data to output_json_path. """ # Read the JSON data from the input file with open(input_json_path, 'r') as f: data = json.load(f) # Loop over every frame and do the replacement for frame in data.get("frames", []): old_path = frame.get("file_path", "") new_path = old_path.replace(old_string, new_string) frame["file_path"] = new_path # Write out the updated JSON data with open(output_json_path, 'w') as f: json.dump(data, f, indent=2) def main(): if len(sys.argv) < 3: print(f"Usage: {sys.argv[0]} ") sys.exit(1) input_json = sys.argv[1] output_json = sys.argv[2] # Call the helper function replace_json_frame_paths(input_json, output_json) if __name__ == "__main__": main() ``` ### data_processing_utils/removepng.py ```python #!/usr/bin/env python3 import json import sys import os def remove_image_extensions_from_json_frames(input_json_path: str, output_json_path: str): """ Reads JSON data from input_json_path, removes .png or .jpg extensions from 'file_path' in the 'frames' list if present, then writes the modified data to output_json_path. """ # Read the JSON data from file with open(input_json_path, 'r') as f: data = json.load(f) # Loop over every frame and remove .png or .jpg from file_path if present for frame in data.get("frames", []): old_path = frame.get("file_path", "") # Use os.path.splitext to split extension; remove only if it's .png or .jpg root, ext = os.path.splitext(old_path) if ext.lower() in [".png", ".jpg"]: frame["file_path"] = root else: frame["file_path"] = old_path # Write out the updated JSON data with open(output_json_path, 'w') as f: json.dump(data, f, indent=2) def main(): if len(sys.argv) < 3: print(f"Usage: {sys.argv[0]} ") sys.exit(1) input_json = sys.argv[1] output_json = sys.argv[2] # Call the helper function remove_image_extensions_from_json_frames(input_json, output_json) if __name__ == "__main__": main() ``` ### data_processing_utils/cleanup.py ```python #!/usr/bin/env python3 import os import json import argparse def remove_missing_frames(input_json_path: str, output_json_path: str, folder_path: str): """ Reads JSON data from input_json_path, removes frames whose corresponding image files do not exist in folder_path, and writes the updated JSON data to output_json_path. """ # 1) Load the JSON data with open(input_json_path, "r") as f: data = json.load(f) frames = data.get("frames", []) print(f"Original frames count: {len(frames)}") # 2) Collect only frames that still exist on disk cleaned_frames = [] for frame in frames: file_path = frame.get("file_path", "") folder_part, filename = os.path.split(file_path) # e.g. "./train", "0007" or "0007.png" # Determine file extension (assume .png if none found) name, ext = os.path.splitext(filename) if not ext: ext = ".png" # Construct the full path to check full_path = os.path.join(folder_path, folder_part, name + ext) full_path = os.path.normpath(full_path) # normalize path if os.path.exists(full_path): cleaned_frames.append(frame) else: print(f"Frame removed - missing file: '{full_path}'") data["frames"] = cleaned_frames print(f"Cleaned frames count: {len(cleaned_frames)}") # 3) Save out the updated JSON with open(output_json_path, "w") as f: json.dump(data, f, indent=2) print(f"Saved cleaned JSON to '{output_json_path}'.") def main(): parser = argparse.ArgumentParser( description="Remove frames from the JSON if their corresponding image file no longer exists." ) parser.add_argument("--input_json", type=str, required=True, help="Path to the original JSON file.") parser.add_argument("--output_json", type=str, required=True, help="Where to write the cleaned JSON.") parser.add_argument("--folder", type=str, required=True, help="Folder where images reside.") args = parser.parse_args() # Call the helper function remove_missing_frames(args.input_json, args.output_json, args.folder) if __name__ == "__main__": main() ``` ### data_processing_utils/scale.py ```python #!/usr/bin/env python3 import os from PIL import Image def scale_images( input_folder: str, output_folder: str, width_factor: float = 0.5, height_factor: float = 0.5 ) -> dict: """ Scales all images in the input_folder by the given width_factor and height_factor, saves them to output_folder, and returns a dictionary mapping each filename to its new dimensions (width, height). :param input_folder: Path to the folder containing the original images. :param output_folder: Path to the folder where scaled images will be saved. :param width_factor: Factor by which to scale the width (e.g., 0.5 for half). :param height_factor: Factor by which to scale the height (e.g., 0.5 for half). :return: A dictionary with filenames as keys and (width, height) tuples as values. """ os.makedirs(output_folder, exist_ok=True) results = {} for filename in os.listdir(input_folder): input_path = os.path.join(input_folder, filename) if not os.path.isfile(input_path): continue try: with Image.open(input_path) as img: original_width, original_height = img.size target_width = int(original_width * width_factor) target_height = int(original_height * height_factor) # Resize using a high-quality downsampling filter. img_resized = img.resize((target_width, target_height), Image.LANCZOS) output_path = os.path.join(output_folder, filename) img_resized.save(output_path) results[filename] = (target_width, target_height) print(f"Scaled '{filename}' -> [{target_width}×{target_height}], saved to '{output_path}'") except Exception as e: print(f"Skipping '{filename}': {e}") return results def main(): # Example usage of the helper function: # Assuming you want to halve both width and height: input_folder = "path_to_input_images" output_folder = "path_to_output_images" # Scale images to half of their original width and height scale_images(input_folder, output_folder, width_factor=0.5, height_factor=0.5) if __name__ == "__main__": main() ``` ### data_processing_utils/split_for_val.py ```python #!/usr/bin/env python3 import os import json import shutil import random import argparse def split_frames_into_val_set( input_json_path: str, output_json_path: str, val_folder: str = "./val", num_val_frames: int = 10, seed: int = 0 ): """ Splits a subset of frames from a transforms JSON file into a validation set, moves the corresponding images into 'val_folder', writes a new validation JSON, and updates the original JSON (overwriting) to exclude those validation frames. :param input_json_path: Path to the original transforms JSON (WILL be overwritten). :param output_json_path: Where to write the new validation JSON (e.g. transforms_val.json). :param val_folder: Folder to move validation images into (must exist or be creatable). :param num_val_frames: How many images to move to the validation set. :param seed: Random seed for reproducibility of frame selection. """ # 1) Read original transforms with open(input_json_path, "r") as f: data = json.load(f) all_frames = data.get("frames", []) total_frames = len(all_frames) if total_frames == 0: print("No frames found in the input JSON. Exiting.") return # 2) Pick random frames for validation random.seed(seed) if num_val_frames > total_frames: print(f"Requested {num_val_frames} val frames, but only {total_frames} exist.") return val_indices = set(random.sample(range(total_frames), num_val_frames)) # 3) Partition frames into 'val_frames' and 'train_frames' val_frames = [] train_frames = [] for i, frame in enumerate(all_frames): if i in val_indices: val_frames.append(frame) else: train_frames.append(frame) os.makedirs(val_folder, exist_ok=True) # 4) Move each chosen val file into the val folder, adjusting the file_path for frame in val_frames: old_path = frame.get("file_path", "") # e.g. "./train/0002" or "./train/0002.png" # We'll guess there's a .png extension if none is present folder_part, filename = os.path.split(old_path) name, ext = os.path.splitext(filename) if ext == "": ext = ".png" # assume .png if missing # The file on disk full_old_path = os.path.join(folder_part, name + ext) # Build new disk path and new JSON path new_disk_path = os.path.join(val_folder, name + ext) new_json_path = f"./val/{name}" # or f"./val/{name}{ext}" if you want to keep .png in JSON # Physically move the file if not os.path.exists(full_old_path): print(f"Warning: image file '{full_old_path}' does not exist!") else: shutil.move(full_old_path, new_disk_path) print(f"Moved '{full_old_path}' -> '{new_disk_path}'") # Update the frame path in the val set frame["file_path"] = new_json_path # 5) Create new JSON data for val val_data = {k: v for k, v in data.items() if k != "frames"} val_data["frames"] = val_frames # 6) Overwrite the original JSON with only the remaining train frames data["frames"] = train_frames with open(input_json_path, "w") as f: json.dump(data, f, indent=2) print(f"Overwrote '{input_json_path}' to keep {len(train_frames)} train frames.") # 7) Write out transforms_val.json with open(output_json_path, "w") as f: json.dump(val_data, f, indent=2) print(f"Wrote '{output_json_path}' with {len(val_frames)} validation frames.") def main(): parser = argparse.ArgumentParser( description="Split a subset of frames from transforms.json into a validation set, " "move images, write transforms_val.json, and update the original transforms." ) parser.add_argument("--input_json", type=str, required=True, help="Path to the original transforms JSON (WILL be overwritten).") parser.add_argument("--output_json", type=str, required=True, help="Where to write the new validation JSON (transforms_val.json).") parser.add_argument("--val_folder", type=str, default="./val", help="Folder to move validation images into (must exist or be creatable).") parser.add_argument("--num_val_frames", type=int, default=10, help="How many images to move to the validation set.") parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility.") args = parser.parse_args() # Call the helper function split_frames_into_val_set( input_json_path=args.input_json, output_json_path=args.output_json, val_folder=args.val_folder, num_val_frames=args.num_val_frames, seed=args.seed ) if __name__ == "__main__": main() ``` ### data_processing_utils/split_for_test.py ```python #!/usr/bin/env python3 import os import json import shutil import random import argparse def split_frames_into_test_set( input_json_path: str, output_json_path: str, test_folder: str = "./test", num_test_frames: int = 10, seed: int = 0 ): """ Splits frames from input_json_path into a test set, moves the corresponding image files into test_folder, writes a new JSON to output_json_path, and updates the original JSON file in-place to remove the test frames. The image file paths in the JSON are interpreted relative to the JSON file's location. :param input_json_path: Path to the original transforms JSON (WILL be overwritten). :param output_json_path: Path to write the new test JSON. :param test_folder: Folder to move test images into (created if needed). :param num_test_frames: How many images to move to the test set. :param seed: Random seed for reproducibility. """ # Determine the base directory of the JSON file. base_dir = os.path.dirname(os.path.abspath(input_json_path)) # 1) Read original transforms with open(input_json_path, "r") as f: data = json.load(f) all_frames = data.get("frames", []) total_frames = len(all_frames) if total_frames == 0: print("No frames found in the input JSON. Exiting.") return # 2) Pick random frames for testing random.seed(seed) if num_test_frames > total_frames: print(f"Requested {num_test_frames} test frames, but only {total_frames} exist.") return test_indices = set(random.sample(range(total_frames), num_test_frames)) # 3) Partition frames into 'test_frames' and 'remaining_frames' test_frames = [] remaining_frames = [] for i, frame in enumerate(all_frames): if i in test_indices: test_frames.append(frame) else: remaining_frames.append(frame) os.makedirs(test_folder, exist_ok=True) # 4) Move each chosen test file into the test folder, adjusting the file_path. for frame in test_frames: rel_path = frame["file_path"] # e.g. "./train/0002" or "./train/0002.png" # Compute the full path relative to the JSON file's location. full_old_path = os.path.join(base_dir, rel_path) # Check if an extension is present; if not, assume .png. folder_part, filename = os.path.split(full_old_path) name, ext = os.path.splitext(filename) if not ext: ext = ".png" full_old_path = os.path.join(folder_part, name + ext) # Build the new disk path and the new JSON file path. new_disk_path = os.path.join(test_folder, name + ext) # Here, we assume the JSON file path should point to the test folder. new_json_path = os.path.join("./" + os.path.basename(test_folder), name) # Move the file physically. if not os.path.exists(full_old_path): print(f"Warning: image file '{full_old_path}' does not exist!") else: shutil.move(full_old_path, new_disk_path) print(f"Moved '{full_old_path}' -> '{new_disk_path}'") # Update the frame's file_path to the new JSON path. frame["file_path"] = new_json_path # 5) Create new JSON data for test. test_data = {k: v for k, v in data.items() if k != "frames"} test_data["frames"] = test_frames # 6) Overwrite the original JSON with only the remaining frames. data["frames"] = remaining_frames with open(input_json_path, "w") as f: json.dump(data, f, indent=2) print(f"Overwrote '{input_json_path}' to keep {len(remaining_frames)} frames (train/val).") # 7) Write out the new test JSON. with open(output_json_path, "w") as f: json.dump(test_data, f, indent=2) print(f"Wrote '{output_json_path}' with {len(test_frames)} test frames.") def main(): parser = argparse.ArgumentParser( description="Split a subset of frames from transforms.json into a test set, " "move images, write transforms_test.json, and update the original transforms." ) parser.add_argument("--input_json", type=str, required=True, help="Path to the original transforms JSON (WILL be overwritten).") parser.add_argument("--output_json", type=str, required=True, help="Where to write the new test JSON (transforms_test.json).") parser.add_argument("--test_folder", type=str, default="./test", help="Folder to move test images into (must exist or be creatable).") parser.add_argument("--num_test_frames", type=int, default=10, help="How many images to move to test set.") parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility.") args = parser.parse_args() split_frames_into_test_set( input_json_path=args.input_json, output_json_path=args.output_json, test_folder=args.test_folder, num_test_frames=args.num_test_frames, seed=args.seed ) if __name__ == "__main__": main() ```