Spaces:
Running
Running
File size: 5,714 Bytes
7088d16 |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import contextlib
import logging
import os
import tempfile
import unittest
from pathlib import Path
from typing import Generator, Tuple
from zipfile import ZipFile
from iopath.common.file_io import PathManager
CO3D_MANIFOLD_PATH: str = "manifold://co3d/tree/extracted"
CO3DV2_MANIFOLD_PATH: str = "manifold://co3d/tree/v2/extracted"
INSIDE_RE_WORKER: bool = os.environ.get("INSIDE_RE_WORKER", False)
def get_path_manager(silence_logs: bool = False) -> PathManager:
"""
Returns a path manager which can access manifold internally.
Args:
silence_logs: Whether to reduce log output from iopath library.
"""
if silence_logs:
logging.getLogger("iopath.fb.manifold").setLevel(logging.CRITICAL)
logging.getLogger("iopath.common.file_io").setLevel(logging.CRITICAL)
if INSIDE_RE_WORKER:
raise ValueError("Cannot get to manifold from RE")
path_manager = PathManager()
if os.environ.get("FB_TEST", False):
from iopath.fb.manifold import ManifoldPathHandler
path_manager.register_handler(ManifoldPathHandler())
return path_manager
@contextlib.contextmanager
def get_skateboard_data(
avoid_manifold: bool = False, silence_logs: bool = False
) -> Generator[Tuple[str, PathManager], None, None]:
"""
Context manager for accessing Co3D dataset by tests, at least for
the first 5 skateboards. Internally, we want this to exercise the
normal way to access the data directly manifold, but on an RE
worker this is impossible so we use a workaround.
Args:
avoid_manifold: Use the method used by RE workers even locally.
silence_logs: Whether to reduce log output from iopath library.
Yields:
dataset_root: (str) path to dataset root.
path_manager: path_manager to access it with.
"""
if silence_logs:
logging.getLogger("iopath.fb.manifold").setLevel(logging.CRITICAL)
logging.getLogger("iopath.common.file_io").setLevel(logging.CRITICAL)
if not os.environ.get("FB_TEST", False):
if os.getenv("FAIR_ENV_CLUSTER", "") == "":
raise unittest.SkipTest("Unknown environment. Data not available.")
yield "/datasets01/co3d/081922", PathManager()
elif avoid_manifold or INSIDE_RE_WORKER:
from libfb.py.parutil import get_file_path
par_path = "skateboard_first_5"
source = get_file_path(par_path)
assert Path(source).is_file()
with tempfile.TemporaryDirectory() as dest:
with ZipFile(source) as f:
f.extractall(dest)
yield os.path.join(dest, "extracted"), PathManager()
else:
yield CO3D_MANIFOLD_PATH, get_path_manager()
def _provide_torchvision_weights(par_path: str, filename: str) -> None:
"""
Ensure the weights files are available for a torchvision model.
"""
# In OSS, torchvision looks for vgg16 weights in
# https://download.pytorch.org/models/vgg16-397923af.pth
# Inside fbcode, this is replaced by asking iopath for
# manifold://torchvision/tree/models/vgg16-397923af.pth
# (the code for this replacement is in
# fbcode/pytorch/vision/fb/_internally_replaced_utils.py )
#
# iopath does this by looking for the file at the cache location
# and if it is not there getting it from manifold.
# (the code for this is in
# fbcode/fair_infra/data/iopath/iopath/fb/manifold.py )
#
# On the remote execution worker, manifold is inaccessible.
# We solve this by making the cached file available before iopath
# looks.
#
# By default the cache location is
# ~/.torch/iopath_cache/manifold_cache/tree/models/vgg16-397923af.pth
# But we can't write to the home directory on the RE worker.
# We define FVCORE_CACHE to change the cache location to
# iopath_cache/manifold_cache/tree/models/vgg16-397923af.pth
# (Without it, manifold caches in unstable temporary locations on RE.)
#
# The file we want has been copied from
# tree/models/vgg16-397923af.pth in the torchvision bucket
# to
# tree/testing/vgg16-397923af.pth in the co3d bucket
# and the TARGETS file copies it somewhere in the PAR which we
# recover with get_file_path.
# (It can't copy straight to a nested location, see
# https://fb.workplace.com/groups/askbuck/posts/2644615728920359/)
# Here we symlink it to the new cache location.
if INSIDE_RE_WORKER:
from libfb.py.parutil import get_file_path
os.environ["FVCORE_CACHE"] = "iopath_cache"
source = Path(get_file_path(par_path))
assert source.is_file()
dest = Path("iopath_cache/manifold_cache/tree/models")
if not dest.exists():
dest.mkdir(parents=True)
if not (dest / filename).is_symlink():
try:
(dest / filename).symlink_to(source)
except FileExistsError:
print("FileExistsError: no symlink created.")
def provide_lpips_vgg() -> None:
"""
Ensure the weights files are available for lpips.LPIPS(net="vgg")
to be called. Specifically, torchvision's vgg16.
"""
_provide_torchvision_weights("vgg_weights_for_lpips", "vgg16-397923af.pth")
def provide_resnet34() -> None:
"""
Ensure the weights files are available for
torchvision.models.resnet34(pretrained=True)
to be called.
"""
_provide_torchvision_weights("resnet34_weights", "resnet34-b627a593.pth")
|