text: import matplotlib.pyplot as plt
  import networkx as nx
  import numpy as np
  import os
  import pickle
  import time
  from tqdm.notebook import tqdm
  
  import torch
  torch.set_default_tensor_type(torch.DoubleTensor)
  
  from spatial_scene_grammars.constraints import *
  from spatial_scene_grammars.nodes import *
  from spatial_scene_grammars.rules import *
  from spatial_scene_grammars.scene_grammar import *
  from spatial_scene_grammars.visualization import *
  from spatial_scene_grammars_examples.dish_bin.grammar import *
  from spatial_scene_grammars.parsing import *
  from spatial_scene_grammars.sampling import *
  from spatial_scene_grammars.parameter_estimation import *
  from spatial_scene_grammars.dataset import *
  
  import meshcat
  import meshcat.geometry as meshcat_geom
  
  # Set up grammar
  grammar = SpatialSceneGrammar(
      root_node_type = DishBin,
      root_node_tf = drake_tf_to_torch_tf(RigidTransform(p=[0.5, 0., 0.]))
  )
  grammar.load_state_dict(torch.load("fit_grammar.torch"))
  
  vis = meshcat.Visualizer()
  vis.delete()
  print("Meshcat port: ", vis.url())
  
  # Convert dataset to observed node sets (caching output) and draw a few examples.
  RECONVERT_DATASET = True
  DATASET_YAML_FILE = "sink/saved_scenes.yaml"
  DATASET_SAVE_FILE = "observed_node_sets.dat"
  
  if RECONVERT_DATASET or not os.path.exists(DATASET_SAVE_FILE):
      type_map = {
          "bin": DishBin
      }
      model_map = {
      }
      for model_type_set in [PlateModels, CupModels, BowlModels]:
          for model_type in model_type_set:
              # Have to cut off the "sink" folder to match model names;
              # dataset management is ugly and should get reorganized...
              model_map[os.path.join(*model_type.sdf.split("/")[1:])] = model_type
      observed_node_sets = convert_scenes_yaml_to_observed_nodes(DATASET_YAML_FILE, type_map, model_map)
      print("Saving...")
      with open(DATASET_SAVE_FILE, "wb") as f:
          pickle.dump(observed_node_sets, f)
  
  print("Loading...")
  with open(DATASET_SAVE_FILE, "rb") as f:
      observed_node_sets = pickle.load(f)
  
  for observed_node_set in observed_node_sets:
      vis.delete()
      draw_scene_tree_contents_meshcat(
          SceneTree.make_from_observed_nodes(observed_node_set),
          zmq_url=vis.window.zmq_url, prefix="scene"
      )
      input()
      