Skip to content

Commit

Permalink
feat: sync up changes from wg-assets
Browse files Browse the repository at this point in the history
- add --directory & --recursive arguments to run over multiple files
- add --camera option to allow for custom cameras
- add --usd-file which is now not required because of --directory
- update docs
  • Loading branch information
beersandrew committed May 2, 2024
1 parent fe0e15e commit 5df0c98
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 29 deletions.
22 changes: 12 additions & 10 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,22 +5,24 @@ Given a USD file, take a picture and assign it as it's thumbnail.

## Usage

`python generate_thumbnail.py <usd_file>`

positional arguments:
- `usd_file` : The USD file you want to add a thumbnail to. If USDZ is input, a new USD file will be created to wrap the existing one called `<subject_usd_file>_Thumbnail.usd`
Single File Generation: `python generate_thumbnail.py --usd-file <usd_file>`
Generation for Entire Directory: `python generate_thumbnail.py --directory <directory>`
Generation for Entire Directory and Recursive Child Folders: `python generate_thumbnail.py --directory <directory> --recursive`

optional arguments:
- `-h`, `--help` : Show help
- `-h`, `--help` : Show help
- `--usd-file` : The USD file you want to add a thumbnail to. If USDZ is input, a new USD file will be created to wrap the existing one called `<subject_usd_file>_Thumbnail.usd`
- `--create-usdz-result` : Returns the resulting files as a new USDZ file called `<subject_usd_file>_Thumbnail.usdz`
- `--output-extension` : The file extension of the output image you want (exr, png..). If using exr, make sure your USD install includes OpenEXR (exr is natively included in USD as of version 23.11)
- `--verbose` : Prints out the steps as they happen
- `--dome-light` : The path to the dome light HDR image to use,
- `--width` : The width of the image in pixels to generate. Default is 2048.

- `--height` : The height of the image in pixels to generate. Default is 2048. If height is not specified, the image is square.
- `--apply-thumbnail`: Saves the image as the thumbnail for the given USD file.
- `--render-purposes`: A comma separated list of render purposes to include in the thumbnail. Valid values are: default, render, proxy, guide.
- `--width` : The width of the image in pixels to generate. Default is 2048.
- `--height` : The height of the image in pixels to generate. Default is 2048. If height is not specified, the image is square.
- `--apply-thumbnail` : Saves the image as the thumbnail for the given USD file.
- `--render-purposes` : A comma separated list of render purposes to include in the thumbnail. Valid values are: default, render, proxy, guide.
- `--directory` : A directory to generate thumbnails for all .usd, .usda, .usdc, and .usdz files. When a directory is supplied, usd-file is ignored.
- `--recursive` : Will recursively search all directories underneath a given directory, requires a directory to be set.
- `--camera` : The path to the camera prim to use for the thumbnail image

Note: You must have usd installed and available in your path. [Install Steps Here](https://github.com/PixarAnimationStudios/OpenUSD#getting-and-building-the-code)

Expand Down
81 changes: 62 additions & 19 deletions generate_thumbnail.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@

def parse_args():
parser = argparse.ArgumentParser(description="This script takes a thumbnail image of the given USD file supplied and associates it with the file.")
parser.add_argument('usd_file',
parser.add_argument('--usd-file',
type=str,
help='The USD file you want to add a thumbnail to. If USDZ is input, a new USD file will be created to wrap the existing one called <input>_Thumbnail.usd')
help='The USD file you want to add a thumbnail to. If USDZ is input, a new USD file will be created to wrap the existing one called `<subject_usd_file>_Thumbnail.usd`')
parser.add_argument('--dome-light',
type=str,
help='The path to the dome light HDR image to use, if any')
Expand Down Expand Up @@ -40,36 +40,48 @@ def parse_args():
type=str,
help='A comma separated list of render purposes to include in the thumbnail. Valid values are: default, render, proxy, guide.',
default='default')

parser.add_argument('--directory',
type=str,
help='A directory to generate thumbnails for all .usd, .usda, .usdc, and .usdz files. When a directory is supplied, usd-file is ignored.')
parser.add_argument('--recursive',
action='store_true',
help='Will recursively search all directories underneath a given directory, requires a directory to be set.')
parser.add_argument('--camera',
type=str,
help='The path to the camera prim to use for the thumbnail image')

return parser.parse_args()

THUMBNAIL_LAYER_SUFFIX_USDA = "_Thumbnail.usda"
THUMBNAIL_LAYER_SUFFIX = "_Thumbnail"
DEFAULT_THUMBNAIL_FILENAME = "default_thumbnail.usda"
DEFAULT_CAMERA_NAME = 'MainCamera'
THUMBNAIL_FOLDER_NAME = "thumbnails"
EXTENSIONS = ('.usd', '.usda', '.usdc', '.usdz')
RENDER_PURPOSE_MAP = {
"default": UsdGeom.Tokens.default_,
"render": UsdGeom.Tokens.render,
"proxy": UsdGeom.Tokens.proxy,
"guide": UsdGeom.Tokens.guide
}

def generate_thumbnail(usd_file, verbose, extension, render_purpose_tokens):
def generate_thumbnail(usd_file, verbose, extension, render_purpose_tokens, camera):
if verbose:
print("Step 1: Setting up the camera...")

subject_stage = Usd.Stage.Open(usd_file)
subject_file = usd_file

setup_camera(subject_stage, subject_file, render_purpose_tokens)

file_to_snapshot = usd_file if camera else get_or_create_file_to_snapshot(subject_stage, subject_file, render_purpose_tokens)
image_path = create_image_filename(usd_file, extension)
camera_to_snapshot_from = camera if camera else DEFAULT_CAMERA_NAME

if verbose:
print("Step 2: Taking the snapshot...")

image_path = create_image_filename(usd_file, extension)
return take_snapshot(image_path)
return take_snapshot(file_to_snapshot, camera_to_snapshot_from, image_path)

def setup_camera(subject_stage, usd_file, render_purpose_tokens):
def get_or_create_file_to_snapshot(subject_stage, usd_file, render_purpose_tokens):
up_axis = UsdGeom.GetStageUpAxis(subject_stage)
is_z_up = up_axis == 'Z'

Expand All @@ -82,6 +94,9 @@ def setup_camera(subject_stage, usd_file, render_purpose_tokens):
sublayer_subject(camera_stage, usd_file)
set_camera_stage_draw_mode(camera_stage, subject_stage)

return DEFAULT_THUMBNAIL_FILENAME


def create_camera(up_axis):
stage = Usd.Stage.CreateNew(DEFAULT_THUMBNAIL_FILENAME)

Expand Down Expand Up @@ -215,11 +230,10 @@ def sublayer_subject(camera_stage, input_file):

return camera_stage

def take_snapshot(image_name):
def take_snapshot(file, camera, image_name):
renderer = get_renderer()
cmd = ['usdrecord', '--camera', 'MainCamera', '--imageWidth', str(args.width), '--renderer', renderer, DEFAULT_THUMBNAIL_FILENAME, image_name]
cmd = ['usdrecord', '--camera', camera, '--imageWidth', str(args.width), '--renderer', renderer, file, image_name]
run_os_specific_command(cmd)
os.remove(DEFAULT_THUMBNAIL_FILENAME)
return image_name

def get_renderer():
Expand Down Expand Up @@ -298,16 +312,16 @@ def list_resolved_dependencies(stage_path):
def convert_render_purposes_to_tokens(render_purposes):
return [RENDER_PURPOSE_MAP[key] for key in render_purposes.split(',')]

if __name__ == "__main__":

args = parse_args()

usd_file = args.usd_file
def generate_single_thumbnail(usd_file, args):
is_usdz = usd_file.endswith(".usdz")

purpose_tokens = convert_render_purposes_to_tokens(args.render_purposes)

image_name = generate_thumbnail(usd_file, args.verbose, args.output_extension, purpose_tokens)
image_name = generate_thumbnail(usd_file, args.verbose, args.output_extension, purpose_tokens, args.camera)

if not args.camera:
os.remove(DEFAULT_THUMBNAIL_FILENAME)

subject_stage = create_usdz_wrapper_stage(usd_file) if is_usdz and args.apply_thumbnail else Usd.Stage.Open(usd_file)

if args.apply_thumbnail:
Expand All @@ -320,4 +334,33 @@ def convert_render_purposes_to_tokens(render_purposes):
if args.verbose:
print("Step 4: Creating usdz result...")

zip_results(usd_file, is_usdz)
zip_results(usd_file, is_usdz)

if __name__ == "__main__":

args = parse_args()

if args.directory and args.recursive:
for root, dirs, files in os.walk(args.directory):
for file in files:
if file.endswith(EXTENSIONS):
file_path = os.path.join(root, file)
print(f"Processing {file_path}...")
try:
generate_single_thumbnail(file_path, args)
except Exception as e:
print(f"Error processing {file_path}: {e}")
elif args.directory:
for file in os.listdir(args.directory):
if file.endswith(EXTENSIONS):
file_path = os.path.join(args.directory, file)
print(f"Processing {file_path}...")
try:
generate_single_thumbnail(file_path, args)
except Exception as e:
print(f"Error processing {file_path}: {e}")
elif args.usd_file:
usd_file = args.usd_file
generate_single_thumbnail(usd_file, args)
else:
print("Either --usd-file or --directory must be set.")

0 comments on commit 5df0c98

Please sign in to comment.