Spaces:
Configuration error
Configuration error
That-Random-Coder
feat: SchwarzNet v2 release - Kerr ray tracer, 5-model deep ensemble, PINN GR verification, Gradio web app
cb40653 | """ | |
| Generate 7000 new images and merge with existing 3000 into a new 10k dataset. | |
| Copies existing data from the old file to avoid re-computing. | |
| """ | |
| import numpy as np | |
| import h5py | |
| import yaml | |
| import os | |
| import shutil | |
| import concurrent.futures | |
| from tqdm import tqdm | |
| from data.generate_dataset import ( | |
| compute_rs_from_solar_mass, | |
| generate_single_image, | |
| add_noise_and_fourier_artifacts, | |
| log1p_transform, | |
| ) | |
| def worker_generate_image(args): | |
| idx, mass_solar, inclination_deg, disk_type, noise_std, fov_factor, config = args | |
| raw_image, rs = generate_single_image(mass_solar, inclination_deg, disk_type, config, fov_factor=fov_factor) | |
| noisy_image = add_noise_and_fourier_artifacts(raw_image, noise_std, float(config['data']['fourier_uv_coverage'])) | |
| normalized_image = log1p_transform(noisy_image) | |
| return idx, normalized_image | |
| def expand_dataset(num_new=7000, config_path='configs/config.yaml'): | |
| with open(config_path, 'r') as f: | |
| config = yaml.safe_load(f) | |
| old_path = config['data']['output_path'] | |
| new_path = old_path.replace('.h5', '_10k.h5') | |
| image_size = int(config['data']['image_size']) | |
| mass_min = float(config['data']['mass_range_solar'][0]) | |
| mass_max = float(config['data']['mass_range_solar'][1]) | |
| incl_min = float(config['data']['inclination_range_deg'][0]) | |
| incl_max = float(config['data']['inclination_range_deg'][1]) | |
| noise_min = float(config['data']['noise_std_range'][0]) | |
| noise_max = float(config['data']['noise_std_range'][1]) | |
| # Read existing data | |
| with h5py.File(old_path, 'r') as hf_old: | |
| existing_count = hf_old['images'].shape[0] | |
| print(f'Existing dataset: {existing_count} images') | |
| total = existing_count + num_new | |
| print(f'Target: {total} images ({num_new} new)') | |
| # Sample parameters for new images | |
| masses_solar = np.exp( | |
| np.random.uniform(np.log(mass_min), np.log(mass_max), num_new) | |
| ).astype(np.float64) | |
| inclinations = np.random.uniform(incl_min, incl_max, num_new).astype(np.float32) | |
| disk_types = np.random.randint(0, 2, num_new) | |
| noise_stds = np.random.uniform(noise_min, noise_max, num_new).astype(np.float32) | |
| fov_factors = np.random.uniform(3.0, 15.0, num_new).astype(np.float32) | |
| rs_values = np.array([ | |
| compute_rs_from_solar_mass(m, config) for m in masses_solar | |
| ], dtype=np.float64) | |
| # Create new file with combined data | |
| with h5py.File(new_path, 'w') as hf_new: | |
| # Create datasets at full size | |
| hf_new.create_dataset('images', shape=(total, 1, image_size, image_size), dtype=np.float32) | |
| hf_new.create_dataset('rs_meters', shape=(total,), dtype=np.float64) | |
| hf_new.create_dataset('mass_solar', shape=(total,), dtype=np.float64) | |
| hf_new.create_dataset('inclination_deg', shape=(total,), dtype=np.float32) | |
| hf_new.create_dataset('disk_type', shape=(total,), dtype=np.int32) | |
| hf_new.create_dataset('noise_std', shape=(total,), dtype=np.float32) | |
| hf_new.create_dataset('fov_factor', shape=(total,), dtype=np.float32) | |
| # Copy existing data | |
| print('Copying existing 3000 images...') | |
| hf_new['images'][:existing_count] = hf_old['images'][:] | |
| hf_new['rs_meters'][:existing_count] = hf_old['rs_meters'][:] | |
| hf_new['mass_solar'][:existing_count] = hf_old['mass_solar'][:] | |
| hf_new['inclination_deg'][:existing_count] = hf_old['inclination_deg'][:] | |
| hf_new['disk_type'][:existing_count] = hf_old['disk_type'][:] | |
| hf_new['noise_std'][:existing_count] = hf_old['noise_std'][:] | |
| hf_new['fov_factor'][:existing_count] = hf_old['fov_factor'][:] | |
| # Write new metadata | |
| hf_new['rs_meters'][existing_count:] = rs_values | |
| hf_new['mass_solar'][existing_count:] = masses_solar | |
| hf_new['inclination_deg'][existing_count:] = inclinations | |
| hf_new['disk_type'][existing_count:] = disk_types | |
| hf_new['noise_std'][existing_count:] = noise_stds | |
| hf_new['fov_factor'][existing_count:] = fov_factors | |
| # Generate new images in parallel | |
| tasks = [ | |
| (idx, masses_solar[idx], inclinations[idx], disk_types[idx], | |
| noise_stds[idx], fov_factors[idx], config) | |
| for idx in range(num_new) | |
| ] | |
| print(f'Generating {num_new} new images in parallel...') | |
| with concurrent.futures.ProcessPoolExecutor() as executor: | |
| futures = {executor.submit(worker_generate_image, task): task[0] for task in tasks} | |
| for future in tqdm(concurrent.futures.as_completed(futures), total=num_new, desc='Generating'): | |
| idx = futures[future] | |
| try: | |
| _, normalized_image = future.result() | |
| hf_new['images'][existing_count + idx, 0] = normalized_image | |
| except Exception as exc: | |
| print(f'Image {idx} generated an exception: {exc}') | |
| # Swap files | |
| backup_path = old_path.replace('.h5', '_3k_backup.h5') | |
| print(f'Backing up old dataset to {backup_path}') | |
| shutil.move(old_path, backup_path) | |
| shutil.move(new_path, old_path) | |
| print(f'\nDone! Dataset expanded: {existing_count} → {total} images') | |
| print(f'Saved to {old_path} (backup at {backup_path})') | |
| if __name__ == '__main__': | |
| expand_dataset() | |