class TestCacheGeometry(unittest.TestCase): def setUp(self): self.cache_dir = random_tempdir('chroma_cache_test') self.cache = Cache(self.cache_dir) self.a = Geometry() self.a.add_solid(Solid(box(1, 1, 1))) self.a.add_solid(Solid(box(1, 1, 1)), displacement=(10, 10, 10)) self.a.flatten() self.b = Geometry() self.b.add_solid(Solid(box(2, 2, 2))) self.b.add_solid(Solid(box(2, 2, 2)), displacement=(10, 10, 10)) self.b.add_solid(Solid(box(2, 2, 2)), displacement=(-10, -10, -10)) self.b.flatten() def test_list_geometry(self): self.assertEqual(len(self.cache.list_geometry()), 0) self.cache.save_geometry('a', self.a) l = self.cache.list_geometry() self.assertEqual(len(l), 1) self.assertIn('a', l) self.cache.save_geometry('b', self.b) l = self.cache.list_geometry() self.assertEqual(len(l), 2) self.assertIn('a', l) self.assertIn('b', l) self.cache.save_geometry('a', self.a) l = self.cache.list_geometry() self.assertEqual(len(l), 2) self.assertIn('a', l) self.assertIn('b', l) def test_load_geometry_not_found(self): with self.assertRaises(GeometryNotFoundError): self.cache.load_geometry('a') def test_save_load_new_geometry(self): self.cache.save_geometry('b', self.b) b = self.cache.load_geometry('b') def test_replace_geometry(self): self.cache.save_geometry('b', self.b) b = self.cache.load_geometry('b') self.assertEqual(b.mesh.md5(), self.b.mesh.md5()) self.cache.save_geometry('b', self.b) b = self.cache.load_geometry('b') self.assertEqual(b.mesh.md5(), self.b.mesh.md5()) def test_remove_geometry(self): self.cache.save_geometry('b', self.b) self.assertIn('b', self.cache.list_geometry()) self.cache.remove_geometry('b') self.assertNotIn('b', self.cache.list_geometry()) def test_get_geometry_hash(self): self.cache.save_geometry('b', self.b) self.assertEqual(self.cache.get_geometry_hash('b'), self.b.mesh.md5()) def test_get_geometry_hash_not_found(self): with self.assertRaises(GeometryNotFoundError): self.cache.get_geometry_hash('a') def test_default_geometry(self): self.cache.save_geometry('a', self.a) self.cache.save_geometry('b', self.b) with self.assertRaises(GeometryNotFoundError): self.cache.set_default_geometry('c') self.cache.set_default_geometry('b') b = self.cache.load_default_geometry() self.cache.set_default_geometry('a') a = self.cache.load_default_geometry() def test_default_geometry_corruption(self): self.cache.save_geometry('a', self.a) self.cache.save_geometry('b', self.b) # Put a file where a symlink should be default_symlink_path = self.cache.get_geometry_filename('.default') with open(default_symlink_path, 'w') as f: f.write('foo') with self.assertRaises(IOError): self.cache.set_default_geometry('b') # Verify file not modified assert os.path.isfile(default_symlink_path) with open(default_symlink_path) as f: self.assertEqual(f.read(), 'foo') def tearDown(self): remove_path(self.cache_dir)
def load_geometry_from_string(geometry_str, auto_build_bvh=True, read_bvh_cache=True, update_bvh_cache=True, cache_dir=None, cuda_device=None): '''Create or load a geometry and optionally load/build a BVH for it. This is a convenience interface to the geometry and BVH construction code, as well as the Chroma caching layer. Most applications should use this function rather than manually building a Geometry and BVH. The geometry string passed to this function has several forms: "" (empty string) - Load the default geometry from the cache and the default BVH for that geometry. "filename.stl" or "filename.stl.bz2" - Create a geometry from a 3D mesh on disk. This model will not be cached, but the BVH can be, depending on whether update_bvh_cache is True. "geometry_name" - Load a geometry from the cache with this name and the default BVH for that geometry. "geometry_name:bvh_name" - Load a geometry from the cache and the requested BVH by name. "@chroma.models.lionsolid" - Run this function inside a Python module, found in the current $PYTHONPATH, to create the geometry, and load the default BVH. For convenience, the current directory is also added to the $PYTHONPATH. "@chroma.models.lionsolid:bvh_name" - Run this function to create the Geometry and load a BVH by name. By default, the Chroma cache in the user's home directory is consulted for both the geometry and the BVH. A different cache directory can be selected by passing the path in via the ``cache_dir`` parameter. If ``read_bvh_cache`` is set to False, then the BVH cache will not be inspected for BVH objects. If the requested BVH (default, or named) does not exist for this geometry (checked by MD5 hashing the geometry mesh) and ``auto_build_bvh`` is true, then a BVH will be automatically generated using the "simple" BVH algorithm. The simple algorithm is very fast, but produces a poor quality BVH. Any newly created BVH will be saved in the Chroma cache if the ``update_cache_bvh`` parameter is True. BVH construction requires a GPU, so the CUDA device number can be specified with the ``cuda_device`` parameter. Returns: a Geometry object (or subclass) with the ``bvh`` property set if the options allow. ''' # Find BVH id if given bvh_name = 'default' if ':' in geometry_str: geometry_id, bvh_name = geometry_str.split(':') else: geometry_id = geometry_str if cache_dir is None: cache = Cache() else: cache = Cache(cache_dir) # Where is the geometry coming from? if os.path.exists(geometry_id) and \ geometry_id.lower().endswith(('.stl', '.bz2')): # Load from file mesh = mesh_from_stl(geometry_id) geometry = Geometry() geometry.add_solid(Solid(mesh, vacuum, vacuum, color=0x33ffffff)) geometry.flatten() elif geometry_id.startswith('@'): # Load from function function_path = geometry_id[1:] module_name, obj_name = function_path.rsplit('.', 1) orig_sys_path = list(sys.path) try: sys.path.append('.') module = __import__(module_name, fromlist=[obj_name]) sys.path = orig_sys_path except ImportError: sys.path = orig_sys_path raise obj = getattr(module, obj_name) geometry = create_geometry_from_obj(obj, bvh_name=bvh_name, auto_build_bvh=auto_build_bvh, read_bvh_cache=read_bvh_cache, update_bvh_cache=update_bvh_cache, cache_dir=cache_dir, cuda_device=cuda_device) return geometry # RETURN EARLY HERE! ALREADY GOT BVH else: # Load from cache if geometry_id == '': geometry = cache.load_default_geometry() else: geometry = cache.load_geometry(geometry_id) # Cached geometries are flattened already geometry.bvh = load_bvh(geometry, bvh_name=bvh_name, auto_build_bvh=auto_build_bvh, read_bvh_cache=read_bvh_cache, update_bvh_cache=update_bvh_cache, cache_dir=cache_dir, cuda_device=cuda_device) return geometry
class TestCacheGeometry(unittest.TestCase): def setUp(self): self.cache_dir = random_tempdir('chroma_cache_test') self.cache = Cache(self.cache_dir) self.a = Geometry() self.a.add_solid(Solid(box(1,1,1))) self.a.add_solid(Solid(box(1,1,1)), displacement=(10,10,10)) self.a.flatten() self.b = Geometry() self.b.add_solid(Solid(box(2,2,2))) self.b.add_solid(Solid(box(2,2,2)), displacement=(10,10,10)) self.b.add_solid(Solid(box(2,2,2)), displacement=(-10,-10,-10)) self.b.flatten() def test_list_geometry(self): self.assertEqual(len(self.cache.list_geometry()), 0) self.cache.save_geometry('a', self.a) l = self.cache.list_geometry() self.assertEqual(len(l), 1) self.assertIn('a', l) self.cache.save_geometry('b', self.b) l = self.cache.list_geometry() self.assertEquals(len(l), 2) self.assertIn('a', l) self.assertIn('b', l) self.cache.save_geometry('a', self.a) l = self.cache.list_geometry() self.assertEquals(len(l), 2) self.assertIn('a', l) self.assertIn('b', l) def test_load_geometry_not_found(self): with self.assertRaises(GeometryNotFoundError): self.cache.load_geometry('a') def test_save_load_new_geometry(self): self.cache.save_geometry('b', self.b) b = self.cache.load_geometry('b') def test_replace_geometry(self): self.cache.save_geometry('b', self.b) b = self.cache.load_geometry('b') self.assertEqual(b.mesh.md5(), self.b.mesh.md5()) self.cache.save_geometry('b', self.b) b = self.cache.load_geometry('b') self.assertEqual(b.mesh.md5(), self.b.mesh.md5()) def test_remove_geometry(self): self.cache.save_geometry('b', self.b) self.assertIn('b', self.cache.list_geometry()) self.cache.remove_geometry('b') self.assertNotIn('b', self.cache.list_geometry()) def test_get_geometry_hash(self): self.cache.save_geometry('b', self.b) self.assertEqual(self.cache.get_geometry_hash('b'), self.b.mesh.md5()) def test_get_geometry_hash_not_found(self): with self.assertRaises(GeometryNotFoundError): self.cache.get_geometry_hash('a') def test_default_geometry(self): self.cache.save_geometry('a', self.a) self.cache.save_geometry('b', self.b) with self.assertRaises(GeometryNotFoundError): self.cache.set_default_geometry('c') self.cache.set_default_geometry('b') b = self.cache.load_default_geometry() self.cache.set_default_geometry('a') a = self.cache.load_default_geometry() def test_default_geometry_corruption(self): self.cache.save_geometry('a', self.a) self.cache.save_geometry('b', self.b) # Put a file where a symlink should be default_symlink_path = self.cache.get_geometry_filename('.default') with open(default_symlink_path, 'w') as f: f.write('foo') with self.assertRaises(IOError): self.cache.set_default_geometry('b') # Verify file not modified assert os.path.isfile(default_symlink_path) with open(default_symlink_path) as f: self.assertEqual(f.read(), 'foo') def tearDown(self): remove_path(self.cache_dir)