File size: 3,103 Bytes
7340df2 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import os
import numpy as np
import open3d as o3d
import trimesh
import random
from tqdm import tqdm
from pathlib import Path
def augment_obj_file(input_path, output_path, n_augmentations=5):
"""
Augment an OBJ file with random transformations
Args:
input_path: Path to input OBJ file
output_path: Directory to save augmented files
n_augmentations: Number of augmented copies to create
"""
# Create output directory if it doesn't exist
os.makedirs(output_path, exist_ok=True)
# Load the original mesh
mesh = trimesh.load(input_path)
original_name = Path(input_path).stem
for i in range(n_augmentations):
# Create a copy of the original mesh
augmented_mesh = mesh.copy()
# Random rotation (0-360 degrees around each axis)
angle_x = np.random.uniform(0, 2*np.pi)
angle_y = np.random.uniform(0, 2*np.pi)
angle_z = np.random.uniform(0, 2*np.pi)
rotation_matrix = trimesh.transformations.euler_matrix(angle_x, angle_y, angle_z)
augmented_mesh.apply_transform(rotation_matrix)
# Random scaling (0.8-1.2 range)
scale_factor = np.random.uniform(0.8, 1.2, size=3)
scale_matrix = np.eye(4)
scale_matrix[:3, :3] *= scale_factor
augmented_mesh.apply_transform(scale_matrix)
# Random translation (-0.1 to 0.1 range in each dimension)
translation = np.random.uniform(-0.1, 0.1, size=3)
translation_matrix = np.eye(4)
translation_matrix[:3, 3] = translation
augmented_mesh.apply_transform(translation_matrix)
# Save the augmented mesh
output_file = os.path.join(output_path, f"{original_name}_aug_{i}.obj")
augmented_mesh.export(output_file)
def augment_all_objs(source_dir, target_dir, n_augmentations=5):
"""
Augment all OBJ files in a directory
Args:
source_dir: Directory containing original OBJ files
target_dir: Directory to save augmented files
n_augmentations: Number of augmented copies per file
"""
# Get all OBJ files in source directory
obj_files = [f for f in os.listdir(source_dir) if f.endswith('.obj')]
print(f"Found {len(obj_files)} OBJ files to augment")
print(f"Will create {n_augmentations} augmented versions per file")
# Process each file
for obj_file in tqdm(obj_files, desc="Augmenting OBJ files"):
input_path = os.path.join(source_dir, obj_file)
augment_obj_file(input_path, target_dir, n_augmentations)
print(f"Finished! Augmented files saved to: {target_dir}")
if __name__ == "__main__":
# Configuration
N_AUGMENTATIONS = 10 # Number of augmented copies per file
SOURCE_DIR = "/root/shapenet_data/train_mesh_data_under_25kb_1000/train" # Directory with original OBJ files
TARGET_DIR = f"/root/shapenet_data/train_mesh_data_under_25kb_1000/train_{N_AUGMENTATIONS}augment" # Where to save augmented files
# Run augmentation
augment_all_objs(SOURCE_DIR, TARGET_DIR, N_AUGMENTATIONS) |