Commit 80a37498 authored by yongshk's avatar yongshk
Browse files

Initial commit

parents
Pipeline #3463 failed with stages
in 0 seconds
# Copyright (c) Facebook, Inc. and its affiliates.
import importlib
import numpy as np
import os
import re
import subprocess
import sys
from collections import defaultdict
import PIL
import torch
import torchvision
from tabulate import tabulate
__all__ = ["collect_env_info"]
def collect_torch_env():
try:
import torch.__config__
return torch.__config__.show()
except ImportError:
# compatible with older versions of pytorch
from torch.utils.collect_env import get_pretty_env_info
return get_pretty_env_info()
def get_env_module():
var_name = "DETECTRON2_ENV_MODULE"
return var_name, os.environ.get(var_name, "<not set>")
def detect_compute_compatibility(CUDA_HOME, so_file):
try:
cuobjdump = os.path.join(CUDA_HOME, "bin", "cuobjdump")
if os.path.isfile(cuobjdump):
output = subprocess.check_output(
"'{}' --list-elf '{}'".format(cuobjdump, so_file), shell=True
)
output = output.decode("utf-8").strip().split("\n")
arch = []
for line in output:
line = re.findall(r"\.sm_([0-9]*)\.", line)[0]
arch.append(".".join(line))
arch = sorted(set(arch))
return ", ".join(arch)
else:
return so_file + "; cannot find cuobjdump"
except Exception:
# unhandled failure
return so_file
def collect_env_info():
has_gpu = torch.cuda.is_available() # true for both CUDA & ROCM
torch_version = torch.__version__
# NOTE that CUDA_HOME/ROCM_HOME could be None even when CUDA runtime libs are functional
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
has_rocm = False
if (getattr(torch.version, "hip", None) is not None) and (ROCM_HOME is not None):
has_rocm = True
has_cuda = has_gpu and (not has_rocm)
data = []
data.append(("sys.platform", sys.platform)) # check-template.yml depends on it
data.append(("Python", sys.version.replace("\n", "")))
data.append(("numpy", np.__version__))
try:
import detectron2 # noqa
data.append(
(
"detectron2",
detectron2.__version__ + " @" + os.path.dirname(detectron2.__file__),
)
)
except ImportError:
data.append(("detectron2", "failed to import"))
except AttributeError:
data.append(("detectron2", "imported a wrong installation"))
try:
import detectron2._C as _C
except ImportError as e:
data.append(("detectron2._C", f"not built correctly: {e}"))
# print system compilers when extension fails to build
if sys.platform != "win32": # don't know what to do for windows
try:
# this is how torch/utils/cpp_extensions.py choose compiler
cxx = os.environ.get("CXX", "c++")
cxx = subprocess.check_output("'{}' --version".format(cxx), shell=True)
cxx = cxx.decode("utf-8").strip().split("\n")[0]
except subprocess.SubprocessError:
cxx = "Not found"
data.append(("Compiler ($CXX)", cxx))
if has_cuda and CUDA_HOME is not None:
try:
nvcc = os.path.join(CUDA_HOME, "bin", "nvcc")
nvcc = subprocess.check_output("'{}' -V".format(nvcc), shell=True)
nvcc = nvcc.decode("utf-8").strip().split("\n")[-1]
except subprocess.SubprocessError:
nvcc = "Not found"
data.append(("CUDA compiler", nvcc))
if has_cuda and sys.platform != "win32":
try:
so_file = importlib.util.find_spec("detectron2._C").origin
except (ImportError, AttributeError):
pass
else:
data.append(
(
"detectron2 arch flags",
detect_compute_compatibility(CUDA_HOME, so_file),
)
)
else:
# print compilers that are used to build extension
data.append(("Compiler", _C.get_compiler_version()))
data.append(("CUDA compiler", _C.get_cuda_version())) # cuda or hip
if has_cuda and getattr(_C, "has_cuda", lambda: True)():
data.append(
(
"detectron2 arch flags",
detect_compute_compatibility(CUDA_HOME, _C.__file__),
)
)
data.append(get_env_module())
data.append(("PyTorch", torch_version + " @" + os.path.dirname(torch.__file__)))
data.append(("PyTorch debug build", torch.version.debug))
try:
data.append(("torch._C._GLIBCXX_USE_CXX11_ABI", torch._C._GLIBCXX_USE_CXX11_ABI))
except Exception:
pass
if not has_gpu:
has_gpu_text = "No: torch.cuda.is_available() == False"
else:
has_gpu_text = "Yes"
data.append(("GPU available", has_gpu_text))
if has_gpu:
devices = defaultdict(list)
for k in range(torch.cuda.device_count()):
cap = ".".join((str(x) for x in torch.cuda.get_device_capability(k)))
name = torch.cuda.get_device_name(k) + f" (arch={cap})"
devices[name].append(str(k))
for name, devids in devices.items():
data.append(("GPU " + ",".join(devids), name))
if has_rocm:
msg = " - invalid!" if not (ROCM_HOME and os.path.isdir(ROCM_HOME)) else ""
data.append(("ROCM_HOME", str(ROCM_HOME) + msg))
else:
try:
from torch.utils.collect_env import (
get_nvidia_driver_version,
run as _run,
)
data.append(("Driver version", get_nvidia_driver_version(_run)))
except Exception:
pass
msg = " - invalid!" if not (CUDA_HOME and os.path.isdir(CUDA_HOME)) else ""
data.append(("CUDA_HOME", str(CUDA_HOME) + msg))
cuda_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
if cuda_arch_list:
data.append(("TORCH_CUDA_ARCH_LIST", cuda_arch_list))
data.append(("Pillow", PIL.__version__))
try:
data.append(
(
"torchvision",
str(torchvision.__version__) + " @" + os.path.dirname(torchvision.__file__),
)
)
if has_cuda:
try:
torchvision_C = importlib.util.find_spec("torchvision._C").origin
msg = detect_compute_compatibility(CUDA_HOME, torchvision_C)
data.append(("torchvision arch flags", msg))
except (ImportError, AttributeError):
data.append(("torchvision._C", "Not found"))
except AttributeError:
data.append(("torchvision", "unknown"))
try:
import fvcore
data.append(("fvcore", fvcore.__version__))
except (ImportError, AttributeError):
pass
try:
import iopath
data.append(("iopath", iopath.__version__))
except (ImportError, AttributeError):
pass
try:
import cv2
data.append(("cv2", cv2.__version__))
except (ImportError, AttributeError):
data.append(("cv2", "Not found"))
env_str = tabulate(data) + "\n"
env_str += collect_torch_env()
return env_str
def test_nccl_ops():
num_gpu = torch.cuda.device_count()
if os.access("/tmp", os.W_OK):
import torch.multiprocessing as mp
dist_url = "file:///tmp/nccl_tmp_file"
print("Testing NCCL connectivity ... this should not hang.")
mp.spawn(_test_nccl_worker, nprocs=num_gpu, args=(num_gpu, dist_url), daemon=False)
print("NCCL succeeded.")
def _test_nccl_worker(rank, num_gpu, dist_url):
import torch.distributed as dist
dist.init_process_group(backend="NCCL", init_method=dist_url, rank=rank, world_size=num_gpu)
dist.barrier(device_ids=[rank])
def main() -> None:
global x
try:
from detectron2.utils.collect_env import collect_env_info as f
print(f())
except ImportError:
print(collect_env_info())
if torch.cuda.is_available():
num_gpu = torch.cuda.device_count()
for k in range(num_gpu):
device = f"cuda:{k}"
try:
x = torch.tensor([1, 2.0], dtype=torch.float32)
x = x.to(device)
except Exception as e:
print(
f"Unable to copy tensor to device={device}: {e}. "
"Your CUDA environment is broken."
)
if num_gpu > 1:
test_nccl_ops()
if __name__ == "__main__":
main() # pragma: no cover
# Copyright (c) Facebook, Inc. and its affiliates.
"""
An awesome colormap for really neat visualizations.
Copied from Detectron, and removed gray colors.
"""
import numpy as np
import random
__all__ = ["colormap", "random_color", "random_colors"]
# fmt: off
# RGB:
_COLORS = np.array(
[
0.000, 0.447, 0.741,
0.850, 0.325, 0.098,
0.929, 0.694, 0.125,
0.494, 0.184, 0.556,
0.466, 0.674, 0.188,
0.301, 0.745, 0.933,
0.635, 0.078, 0.184,
0.300, 0.300, 0.300,
0.600, 0.600, 0.600,
1.000, 0.000, 0.000,
1.000, 0.500, 0.000,
0.749, 0.749, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 1.000,
0.667, 0.000, 1.000,
0.333, 0.333, 0.000,
0.333, 0.667, 0.000,
0.333, 1.000, 0.000,
0.667, 0.333, 0.000,
0.667, 0.667, 0.000,
0.667, 1.000, 0.000,
1.000, 0.333, 0.000,
1.000, 0.667, 0.000,
1.000, 1.000, 0.000,
0.000, 0.333, 0.500,
0.000, 0.667, 0.500,
0.000, 1.000, 0.500,
0.333, 0.000, 0.500,
0.333, 0.333, 0.500,
0.333, 0.667, 0.500,
0.333, 1.000, 0.500,
0.667, 0.000, 0.500,
0.667, 0.333, 0.500,
0.667, 0.667, 0.500,
0.667, 1.000, 0.500,
1.000, 0.000, 0.500,
1.000, 0.333, 0.500,
1.000, 0.667, 0.500,
1.000, 1.000, 0.500,
0.000, 0.333, 1.000,
0.000, 0.667, 1.000,
0.000, 1.000, 1.000,
0.333, 0.000, 1.000,
0.333, 0.333, 1.000,
0.333, 0.667, 1.000,
0.333, 1.000, 1.000,
0.667, 0.000, 1.000,
0.667, 0.333, 1.000,
0.667, 0.667, 1.000,
0.667, 1.000, 1.000,
1.000, 0.000, 1.000,
1.000, 0.333, 1.000,
1.000, 0.667, 1.000,
0.333, 0.000, 0.000,
0.500, 0.000, 0.000,
0.667, 0.000, 0.000,
0.833, 0.000, 0.000,
1.000, 0.000, 0.000,
0.000, 0.167, 0.000,
0.000, 0.333, 0.000,
0.000, 0.500, 0.000,
0.000, 0.667, 0.000,
0.000, 0.833, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 0.167,
0.000, 0.000, 0.333,
0.000, 0.000, 0.500,
0.000, 0.000, 0.667,
0.000, 0.000, 0.833,
0.000, 0.000, 1.000,
0.000, 0.000, 0.000,
0.143, 0.143, 0.143,
0.857, 0.857, 0.857,
1.000, 1.000, 1.000
]
).astype(np.float32).reshape(-1, 3)
# fmt: on
def colormap(rgb=False, maximum=255):
"""
Args:
rgb (bool): whether to return RGB colors or BGR colors.
maximum (int): either 255 or 1
Returns:
ndarray: a float32 array of Nx3 colors, in range [0, 255] or [0, 1]
"""
assert maximum in [255, 1], maximum
c = _COLORS * maximum
if not rgb:
c = c[:, ::-1]
return c
def random_color(rgb=False, maximum=255):
"""
Args:
rgb (bool): whether to return RGB colors or BGR colors.
maximum (int): either 255 or 1
Returns:
ndarray: a vector of 3 numbers
"""
idx = np.random.randint(0, len(_COLORS))
ret = _COLORS[idx] * maximum
if not rgb:
ret = ret[::-1]
return ret
def random_colors(N, rgb=False, maximum=255):
"""
Args:
N (int): number of unique colors needed
rgb (bool): whether to return RGB colors or BGR colors.
maximum (int): either 255 or 1
Returns:
ndarray: a list of random_color
"""
indices = random.sample(range(len(_COLORS)), N)
ret = [_COLORS[i] * maximum for i in indices]
if not rgb:
ret = [x[::-1] for x in ret]
return ret
if __name__ == "__main__":
import cv2
size = 100
H, W = 10, 10
canvas = np.random.rand(H * size, W * size, 3).astype("float32")
for h in range(H):
for w in range(W):
idx = h * W + w
if idx >= len(_COLORS):
break
canvas[h * size : (h + 1) * size, w * size : (w + 1) * size] = _COLORS[idx]
cv2.imshow("a", canvas)
cv2.waitKey(0)
# Copyright (c) Facebook, Inc. and its affiliates.
"""
This file contains primitives for multi-gpu communication.
This is useful when doing distributed training.
"""
import functools
import numpy as np
import torch
import torch.distributed as dist
_LOCAL_PROCESS_GROUP = None
_MISSING_LOCAL_PG_ERROR = (
"Local process group is not yet created! Please use detectron2's `launch()` "
"to start processes and initialize pytorch process group. If you need to start "
"processes in other ways, please call comm.create_local_process_group("
"num_workers_per_machine) after calling torch.distributed.init_process_group()."
)
def get_world_size() -> int:
if not dist.is_available():
return 1
if not dist.is_initialized():
return 1
return dist.get_world_size()
def get_rank() -> int:
if not dist.is_available():
return 0
if not dist.is_initialized():
return 0
return dist.get_rank()
@functools.lru_cache()
def create_local_process_group(num_workers_per_machine: int) -> None:
"""
Create a process group that contains ranks within the same machine.
Detectron2's launch() in engine/launch.py will call this function. If you start
workers without launch(), you'll have to also call this. Otherwise utilities
like `get_local_rank()` will not work.
This function contains a barrier. All processes must call it together.
Args:
num_workers_per_machine: the number of worker processes per machine. Typically
the number of GPUs.
"""
global _LOCAL_PROCESS_GROUP
assert _LOCAL_PROCESS_GROUP is None
assert get_world_size() % num_workers_per_machine == 0
num_machines = get_world_size() // num_workers_per_machine
machine_rank = get_rank() // num_workers_per_machine
for i in range(num_machines):
ranks_on_i = list(range(i * num_workers_per_machine, (i + 1) * num_workers_per_machine))
pg = dist.new_group(ranks_on_i)
if i == machine_rank:
_LOCAL_PROCESS_GROUP = pg
def get_local_process_group():
"""
Returns:
A torch process group which only includes processes that are on the same
machine as the current process. This group can be useful for communication
within a machine, e.g. a per-machine SyncBN.
"""
assert _LOCAL_PROCESS_GROUP is not None, _MISSING_LOCAL_PG_ERROR
return _LOCAL_PROCESS_GROUP
def get_local_rank() -> int:
"""
Returns:
The rank of the current process within the local (per-machine) process group.
"""
if not dist.is_available():
return 0
if not dist.is_initialized():
return 0
assert _LOCAL_PROCESS_GROUP is not None, _MISSING_LOCAL_PG_ERROR
return dist.get_rank(group=_LOCAL_PROCESS_GROUP)
def get_local_size() -> int:
"""
Returns:
The size of the per-machine process group,
i.e. the number of processes per machine.
"""
if not dist.is_available():
return 1
if not dist.is_initialized():
return 1
assert _LOCAL_PROCESS_GROUP is not None, _MISSING_LOCAL_PG_ERROR
return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)
def is_main_process() -> bool:
return get_rank() == 0
def synchronize():
"""
Helper function to synchronize (barrier) among all processes when
using distributed training
"""
if not dist.is_available():
return
if not dist.is_initialized():
return
world_size = dist.get_world_size()
if world_size == 1:
return
if dist.get_backend() == dist.Backend.NCCL:
# This argument is needed to avoid warnings.
# It's valid only for NCCL backend.
dist.barrier(device_ids=[torch.cuda.current_device()])
else:
dist.barrier()
@functools.lru_cache()
def _get_global_gloo_group():
"""
Return a process group based on gloo backend, containing all the ranks
The result is cached.
"""
if dist.get_backend() == "nccl":
return dist.new_group(backend="gloo")
else:
return dist.group.WORLD
def all_gather(data, group=None):
"""
Run all_gather on arbitrary picklable data (not necessarily tensors).
Args:
data: any picklable object
group: a torch process group. By default, will use a group which
contains all ranks on gloo backend.
Returns:
list[data]: list of data gathered from each rank
"""
if get_world_size() == 1:
return [data]
if group is None:
group = _get_global_gloo_group() # use CPU group by default, to reduce GPU RAM usage.
world_size = dist.get_world_size(group)
if world_size == 1:
return [data]
output = [None for _ in range(world_size)]
dist.all_gather_object(output, data, group=group)
return output
def gather(data, dst=0, group=None):
"""
Run gather on arbitrary picklable data (not necessarily tensors).
Args:
data: any picklable object
dst (int): destination rank
group: a torch process group. By default, will use a group which
contains all ranks on gloo backend.
Returns:
list[data]: on dst, a list of data gathered from each rank. Otherwise,
an empty list.
"""
if get_world_size() == 1:
return [data]
if group is None:
group = _get_global_gloo_group()
world_size = dist.get_world_size(group=group)
if world_size == 1:
return [data]
rank = dist.get_rank(group=group)
if rank == dst:
output = [None for _ in range(world_size)]
dist.gather_object(data, output, dst=dst, group=group)
return output
else:
dist.gather_object(data, None, dst=dst, group=group)
return []
def shared_random_seed():
"""
Returns:
int: a random number that is the same across all workers.
If workers need a shared RNG, they can use this shared seed to
create one.
All workers must call this function, otherwise it will deadlock.
"""
ints = np.random.randint(2**31)
all_ints = all_gather(ints)
return all_ints[0]
def reduce_dict(input_dict, average=True):
"""
Reduce the values in the dictionary from all processes so that process with rank
0 has the reduced results.
Args:
input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
average (bool): whether to do average or sum
Returns:
a dict with the same keys as input_dict, after reduction.
"""
world_size = get_world_size()
if world_size < 2:
return input_dict
with torch.no_grad():
names = []
values = []
# sort the keys so that they are consistent across processes
for k in sorted(input_dict.keys()):
names.append(k)
values.append(input_dict[k])
values = torch.stack(values, dim=0)
dist.reduce(values, dst=0)
if dist.get_rank() == 0 and average:
# only main process gets accumulated, so only divide by
# world_size in this case
values /= world_size
reduced_dict = {k: v for k, v in zip(names, values)}
return reduced_dict
# Copyright (c) Facebook, Inc. and its affiliates.
""" Utilities for developers only.
These are not visible to users (not automatically imported). And should not
appeared in docs."""
# adapted from https://github.com/tensorpack/tensorpack/blob/master/tensorpack/utils/develop.py
def create_dummy_class(klass, dependency, message=""):
"""
When a dependency of a class is not available, create a dummy class which throws ImportError
when used.
Args:
klass (str): name of the class.
dependency (str): name of the dependency.
message: extra message to print
Returns:
class: a class object
"""
err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, klass)
if message:
err = err + " " + message
class _DummyMetaClass(type):
# throw error on class attribute access
def __getattr__(_, __): # noqa: B902
raise ImportError(err)
class _Dummy(object, metaclass=_DummyMetaClass):
# throw error on constructor
def __init__(self, *args, **kwargs):
raise ImportError(err)
return _Dummy
def create_dummy_func(func, dependency, message=""):
"""
When a dependency of a function is not available, create a dummy function which throws
ImportError when used.
Args:
func (str): name of the function.
dependency (str or list[str]): name(s) of the dependency.
message: extra message to print
Returns:
function: a function object
"""
err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, func)
if message:
err = err + " " + message
if isinstance(dependency, (list, tuple)):
dependency = ",".join(dependency)
def _dummy(*args, **kwargs):
raise ImportError(err)
return _dummy
# Copyright (c) Facebook, Inc. and its affiliates.
import importlib
import importlib.util
import logging
import numpy as np
import os
import random
import sys
from datetime import datetime
import torch
__all__ = ["seed_all_rng"]
TORCH_VERSION = tuple(int(x) for x in torch.__version__.split(".")[:2])
"""
PyTorch version as a tuple of 2 ints. Useful for comparison.
"""
DOC_BUILDING = os.getenv("_DOC_BUILDING", False) # set in docs/conf.py
"""
Whether we're building documentation.
"""
def seed_all_rng(seed=None):
"""
Set the random seed for the RNG in torch, numpy and python.
Args:
seed (int): if None, will use a strong random seed.
"""
if seed is None:
seed = (
os.getpid()
+ int(datetime.now().strftime("%S%f"))
+ int.from_bytes(os.urandom(2), "big")
)
logger = logging.getLogger(__name__)
logger.info("Using a generated random seed {}".format(seed))
np.random.seed(seed)
torch.manual_seed(seed)
random.seed(seed)
torch.cuda.manual_seed_all(str(seed))
os.environ["PYTHONHASHSEED"] = str(seed)
# from https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path
def _import_file(module_name, file_path, make_importable=False):
spec = importlib.util.spec_from_file_location(module_name, file_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
if make_importable:
sys.modules[module_name] = module
return module
def _configure_libraries():
"""
Configurations for some libraries.
"""
# An environment option to disable `import cv2` globally,
# in case it leads to negative performance impact
disable_cv2 = int(os.environ.get("DETECTRON2_DISABLE_CV2", False))
if disable_cv2:
sys.modules["cv2"] = None
else:
# Disable opencl in opencv since its interaction with cuda often has negative effects
# This envvar is supported after OpenCV 3.4.0
os.environ["OPENCV_OPENCL_RUNTIME"] = "disabled"
try:
import cv2
if int(cv2.__version__.split(".")[0]) >= 3:
cv2.ocl.setUseOpenCL(False)
except ModuleNotFoundError:
# Other types of ImportError, if happened, should not be ignored.
# Because a failed opencv import could mess up address space
# https://github.com/skvark/opencv-python/issues/381
pass
def get_version(module, digit=2):
return tuple(map(int, module.__version__.split(".")[:digit]))
# fmt: off
assert get_version(torch) >= (1, 4), "Requires torch>=1.4"
import fvcore
assert get_version(fvcore, 3) >= (0, 1, 2), "Requires fvcore>=0.1.2"
import yaml
assert get_version(yaml) >= (5, 1), "Requires pyyaml>=5.1"
# fmt: on
_ENV_SETUP_DONE = False
def setup_environment():
"""Perform environment setup work. The default setup is a no-op, but this
function allows the user to specify a Python source file or a module in
the $DETECTRON2_ENV_MODULE environment variable, that performs
custom setup work that may be necessary to their computing environment.
"""
global _ENV_SETUP_DONE
if _ENV_SETUP_DONE:
return
_ENV_SETUP_DONE = True
_configure_libraries()
custom_module_path = os.environ.get("DETECTRON2_ENV_MODULE")
if custom_module_path:
setup_custom_environment(custom_module_path)
else:
# The default setup is a no-op
pass
def setup_custom_environment(custom_module):
"""
Load custom environment setup by importing a Python source file or a
module, and run the setup function.
"""
if custom_module.endswith(".py"):
module = _import_file("detectron2.utils.env.custom_module", custom_module)
else:
module = importlib.import_module(custom_module)
assert hasattr(module, "setup_environment") and callable(module.setup_environment), (
"Custom environment module defined in {} does not have the "
"required callable attribute 'setup_environment'."
).format(custom_module)
module.setup_environment()
def fixup_module_metadata(module_name, namespace, keys=None):
"""
Fix the __qualname__ of module members to be their exported api name, so
when they are referenced in docs, sphinx can find them. Reference:
https://github.com/python-trio/trio/blob/6754c74eacfad9cc5c92d5c24727a2f3b620624e/trio/_util.py#L216-L241
"""
if not DOC_BUILDING:
return
seen_ids = set()
def fix_one(qualname, name, obj):
# avoid infinite recursion (relevant when using
# typing.Generic, for example)
if id(obj) in seen_ids:
return
seen_ids.add(id(obj))
mod = getattr(obj, "__module__", None)
if mod is not None and (mod.startswith(module_name) or mod.startswith("fvcore.")):
obj.__module__ = module_name
# Modules, unlike everything else in Python, put fully-qualitied
# names into their __name__ attribute. We check for "." to avoid
# rewriting these.
if hasattr(obj, "__name__") and "." not in obj.__name__:
obj.__name__ = name
obj.__qualname__ = qualname
if isinstance(obj, type):
for attr_name, attr_value in obj.__dict__.items():
fix_one(objname + "." + attr_name, attr_name, attr_value)
if keys is None:
keys = namespace.keys()
for objname in keys:
if not objname.startswith("_"):
obj = namespace[objname]
fix_one(objname, objname, obj)
# Copyright (c) Facebook, Inc. and its affiliates.
import datetime
import json
import logging
import os
import time
from collections import defaultdict
from contextlib import contextmanager
from functools import cached_property
from typing import Optional
import torch
from fvcore.common.history_buffer import HistoryBuffer
from detectron2.utils.file_io import PathManager
__all__ = [
"get_event_storage",
"has_event_storage",
"JSONWriter",
"TensorboardXWriter",
"CommonMetricPrinter",
"EventStorage",
]
_CURRENT_STORAGE_STACK = []
def get_event_storage():
"""
Returns:
The :class:`EventStorage` object that's currently being used.
Throws an error if no :class:`EventStorage` is currently enabled.
"""
assert len(
_CURRENT_STORAGE_STACK
), "get_event_storage() has to be called inside a 'with EventStorage(...)' context!"
return _CURRENT_STORAGE_STACK[-1]
def has_event_storage():
"""
Returns:
Check if there are EventStorage() context existed.
"""
return len(_CURRENT_STORAGE_STACK) > 0
class EventWriter:
"""
Base class for writers that obtain events from :class:`EventStorage` and process them.
"""
def write(self):
raise NotImplementedError
def close(self):
pass
class JSONWriter(EventWriter):
"""
Write scalars to a json file.
It saves scalars as one json per line (instead of a big json) for easy parsing.
Examples parsing such a json file:
::
$ cat metrics.json | jq -s '.[0:2]'
[
{
"data_time": 0.008433341979980469,
"iteration": 19,
"loss": 1.9228371381759644,
"loss_box_reg": 0.050025828182697296,
"loss_classifier": 0.5316952466964722,
"loss_mask": 0.7236229181289673,
"loss_rpn_box": 0.0856662318110466,
"loss_rpn_cls": 0.48198649287223816,
"lr": 0.007173333333333333,
"time": 0.25401854515075684
},
{
"data_time": 0.007216215133666992,
"iteration": 39,
"loss": 1.282649278640747,
"loss_box_reg": 0.06222952902317047,
"loss_classifier": 0.30682939291000366,
"loss_mask": 0.6970193982124329,
"loss_rpn_box": 0.038663312792778015,
"loss_rpn_cls": 0.1471673548221588,
"lr": 0.007706666666666667,
"time": 0.2490077018737793
}
]
$ cat metrics.json | jq '.loss_mask'
0.7126231789588928
0.689423680305481
0.6776131987571716
...
"""
def __init__(self, json_file, window_size=20):
"""
Args:
json_file (str): path to the json file. New data will be appended if the file exists.
window_size (int): the window size of median smoothing for the scalars whose
`smoothing_hint` are True.
"""
self._file_handle = PathManager.open(json_file, "a")
self._window_size = window_size
self._last_write = -1
def write(self):
storage = get_event_storage()
to_save = defaultdict(dict)
for k, (v, iter) in storage.latest_with_smoothing_hint(self._window_size).items():
# keep scalars that have not been written
if iter <= self._last_write:
continue
to_save[iter][k] = v
if len(to_save):
all_iters = sorted(to_save.keys())
self._last_write = max(all_iters)
for itr, scalars_per_iter in to_save.items():
scalars_per_iter["iteration"] = itr
self._file_handle.write(json.dumps(scalars_per_iter, sort_keys=True) + "\n")
self._file_handle.flush()
try:
os.fsync(self._file_handle.fileno())
except AttributeError:
pass
def close(self):
self._file_handle.close()
class TensorboardXWriter(EventWriter):
"""
Write all scalars to a tensorboard file.
"""
def __init__(self, log_dir: str, window_size: int = 20, **kwargs):
"""
Args:
log_dir (str): the directory to save the output events
window_size (int): the scalars will be median-smoothed by this window size
kwargs: other arguments passed to `torch.utils.tensorboard.SummaryWriter(...)`
"""
self._window_size = window_size
self._writer_args = {"log_dir": log_dir, **kwargs}
self._last_write = -1
@cached_property
def _writer(self):
from torch.utils.tensorboard import SummaryWriter
return SummaryWriter(**self._writer_args)
def write(self):
storage = get_event_storage()
new_last_write = self._last_write
for k, (v, iter) in storage.latest_with_smoothing_hint(self._window_size).items():
if iter > self._last_write:
self._writer.add_scalar(k, v, iter)
new_last_write = max(new_last_write, iter)
self._last_write = new_last_write
# storage.put_{image,histogram} is only meant to be used by
# tensorboard writer. So we access its internal fields directly from here.
if len(storage._vis_data) >= 1:
for img_name, img, step_num in storage._vis_data:
self._writer.add_image(img_name, img, step_num)
# Storage stores all image data and rely on this writer to clear them.
# As a result it assumes only one writer will use its image data.
# An alternative design is to let storage store limited recent
# data (e.g. only the most recent image) that all writers can access.
# In that case a writer may not see all image data if its period is long.
storage.clear_images()
if len(storage._histograms) >= 1:
for params in storage._histograms:
self._writer.add_histogram_raw(**params)
storage.clear_histograms()
def close(self):
if "_writer" in self.__dict__:
self._writer.close()
class CommonMetricPrinter(EventWriter):
"""
Print **common** metrics to the terminal, including
iteration time, ETA, memory, all losses, and the learning rate.
It also applies smoothing using a window of 20 elements.
It's meant to print common metrics in common ways.
To print something in more customized ways, please implement a similar printer by yourself.
"""
def __init__(self, max_iter: Optional[int] = None, window_size: int = 20):
"""
Args:
max_iter: the maximum number of iterations to train.
Used to compute ETA. If not given, ETA will not be printed.
window_size (int): the losses will be median-smoothed by this window size
"""
self.logger = logging.getLogger("detectron2.utils.events")
self._max_iter = max_iter
self._window_size = window_size
self._last_write = None # (step, time) of last call to write(). Used to compute ETA
def _get_eta(self, storage) -> Optional[str]:
if self._max_iter is None:
return ""
iteration = storage.iter
try:
eta_seconds = storage.history("time").median(1000) * (self._max_iter - iteration - 1)
storage.put_scalar("eta_seconds", eta_seconds, smoothing_hint=False)
return str(datetime.timedelta(seconds=int(eta_seconds)))
except KeyError:
# estimate eta on our own - more noisy
eta_string = None
if self._last_write is not None:
estimate_iter_time = (time.perf_counter() - self._last_write[1]) / (
iteration - self._last_write[0]
)
eta_seconds = estimate_iter_time * (self._max_iter - iteration - 1)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
self._last_write = (iteration, time.perf_counter())
return eta_string
def write(self):
storage = get_event_storage()
iteration = storage.iter
if iteration == self._max_iter:
# This hook only reports training progress (loss, ETA, etc) but not other data,
# therefore do not write anything after training succeeds, even if this method
# is called.
return
try:
avg_data_time = storage.history("data_time").avg(
storage.count_samples("data_time", self._window_size)
)
last_data_time = storage.history("data_time").latest()
except KeyError:
# they may not exist in the first few iterations (due to warmup)
# or when SimpleTrainer is not used
avg_data_time = None
last_data_time = None
try:
avg_iter_time = storage.history("time").global_avg()
last_iter_time = storage.history("time").latest()
except KeyError:
avg_iter_time = None
last_iter_time = None
try:
lr = "{:.5g}".format(storage.history("lr").latest())
except KeyError:
lr = "N/A"
eta_string = self._get_eta(storage)
if torch.cuda.is_available():
max_mem_mb = torch.cuda.max_memory_allocated() / 1024.0 / 1024.0
else:
max_mem_mb = None
# NOTE: max_mem is parsed by grep in "dev/parse_results.sh"
self.logger.info(
str.format(
" {eta}iter: {iter} {losses} {non_losses} {avg_time}{last_time}"
+ "{avg_data_time}{last_data_time} lr: {lr} {memory}",
eta=f"eta: {eta_string} " if eta_string else "",
iter=iteration,
losses=" ".join(
[
"{}: {:.4g}".format(
k, v.median(storage.count_samples(k, self._window_size))
)
for k, v in storage.histories().items()
if "loss" in k
]
),
non_losses=" ".join(
[
"{}: {:.4g}".format(
k, v.median(storage.count_samples(k, self._window_size))
)
for k, v in storage.histories().items()
if "[metric]" in k
]
),
avg_time=(
"time: {:.4f} ".format(avg_iter_time) if avg_iter_time is not None else ""
),
last_time=(
"last_time: {:.4f} ".format(last_iter_time)
if last_iter_time is not None
else ""
),
avg_data_time=(
"data_time: {:.4f} ".format(avg_data_time) if avg_data_time is not None else ""
),
last_data_time=(
"last_data_time: {:.4f} ".format(last_data_time)
if last_data_time is not None
else ""
),
lr=lr,
memory="max_mem: {:.0f}M".format(max_mem_mb) if max_mem_mb is not None else "",
)
)
class EventStorage:
"""
The user-facing class that provides metric storage functionalities.
In the future we may add support for storing / logging other types of data if needed.
"""
def __init__(self, start_iter=0):
"""
Args:
start_iter (int): the iteration number to start with
"""
self._history = defaultdict(HistoryBuffer)
self._smoothing_hints = {}
self._latest_scalars = {}
self._iter = start_iter
self._current_prefix = ""
self._vis_data = []
self._histograms = []
def put_image(self, img_name, img_tensor):
"""
Add an `img_tensor` associated with `img_name`, to be shown on
tensorboard.
Args:
img_name (str): The name of the image to put into tensorboard.
img_tensor (torch.Tensor or numpy.array): An `uint8` or `float`
Tensor of shape `[channel, height, width]` where `channel` is
3. The image format should be RGB. The elements in img_tensor
can either have values in [0, 1] (float32) or [0, 255] (uint8).
The `img_tensor` will be visualized in tensorboard.
"""
self._vis_data.append((img_name, img_tensor, self._iter))
def put_scalar(self, name, value, smoothing_hint=True, cur_iter=None):
"""
Add a scalar `value` to the `HistoryBuffer` associated with `name`.
Args:
smoothing_hint (bool): a 'hint' on whether this scalar is noisy and should be
smoothed when logged. The hint will be accessible through
:meth:`EventStorage.smoothing_hints`. A writer may ignore the hint
and apply custom smoothing rule.
It defaults to True because most scalars we save need to be smoothed to
provide any useful signal.
cur_iter (int): an iteration number to set explicitly instead of current iteration
"""
name = self._current_prefix + name
cur_iter = self._iter if cur_iter is None else cur_iter
history = self._history[name]
value = float(value)
history.update(value, cur_iter)
self._latest_scalars[name] = (value, cur_iter)
existing_hint = self._smoothing_hints.get(name)
if existing_hint is not None:
assert (
existing_hint == smoothing_hint
), "Scalar {} was put with a different smoothing_hint!".format(name)
else:
self._smoothing_hints[name] = smoothing_hint
def put_scalars(self, *, smoothing_hint=True, cur_iter=None, **kwargs):
"""
Put multiple scalars from keyword arguments.
Examples:
storage.put_scalars(loss=my_loss, accuracy=my_accuracy, smoothing_hint=True)
"""
for k, v in kwargs.items():
self.put_scalar(k, v, smoothing_hint=smoothing_hint, cur_iter=cur_iter)
def put_histogram(self, hist_name, hist_tensor, bins=1000):
"""
Create a histogram from a tensor.
Args:
hist_name (str): The name of the histogram to put into tensorboard.
hist_tensor (torch.Tensor): A Tensor of arbitrary shape to be converted
into a histogram.
bins (int): Number of histogram bins.
"""
ht_min, ht_max = hist_tensor.min().item(), hist_tensor.max().item()
# Create a histogram with PyTorch
hist_counts = torch.histc(hist_tensor, bins=bins)
hist_edges = torch.linspace(start=ht_min, end=ht_max, steps=bins + 1, dtype=torch.float32)
# Parameter for the add_histogram_raw function of SummaryWriter
hist_params = dict(
tag=hist_name,
min=ht_min,
max=ht_max,
num=len(hist_tensor),
sum=float(hist_tensor.sum()),
sum_squares=float(torch.sum(hist_tensor**2)),
bucket_limits=hist_edges[1:].tolist(),
bucket_counts=hist_counts.tolist(),
global_step=self._iter,
)
self._histograms.append(hist_params)
def history(self, name):
"""
Returns:
HistoryBuffer: the scalar history for name
"""
ret = self._history.get(name, None)
if ret is None:
raise KeyError("No history metric available for {}!".format(name))
return ret
def histories(self):
"""
Returns:
dict[name -> HistoryBuffer]: the HistoryBuffer for all scalars
"""
return self._history
def latest(self):
"""
Returns:
dict[str -> (float, int)]: mapping from the name of each scalar to the most
recent value and the iteration number its added.
"""
return self._latest_scalars
def latest_with_smoothing_hint(self, window_size=20):
"""
Similar to :meth:`latest`, but the returned values
are either the un-smoothed original latest value,
or a median of the given window_size,
depend on whether the smoothing_hint is True.
This provides a default behavior that other writers can use.
Note: All scalars saved in the past `window_size` iterations are used for smoothing.
This is different from the `window_size` definition in HistoryBuffer.
Use :meth:`get_history_window_size` to get the `window_size` used in HistoryBuffer.
"""
result = {}
for k, (v, itr) in self._latest_scalars.items():
result[k] = (
(
self._history[k].median(self.count_samples(k, window_size))
if self._smoothing_hints[k]
else v
),
itr,
)
return result
def count_samples(self, name, window_size=20):
"""
Return the number of samples logged in the past `window_size` iterations.
"""
samples = 0
data = self._history[name].values()
for _, iter_ in reversed(data):
if iter_ > data[-1][1] - window_size:
samples += 1
else:
break
return samples
def smoothing_hints(self):
"""
Returns:
dict[name -> bool]: the user-provided hint on whether the scalar
is noisy and needs smoothing.
"""
return self._smoothing_hints
def step(self):
"""
User should either: (1) Call this function to increment storage.iter when needed. Or
(2) Set `storage.iter` to the correct iteration number before each iteration.
The storage will then be able to associate the new data with an iteration number.
"""
self._iter += 1
@property
def iter(self):
"""
Returns:
int: The current iteration number. When used together with a trainer,
this is ensured to be the same as trainer.iter.
"""
return self._iter
@iter.setter
def iter(self, val):
self._iter = int(val)
@property
def iteration(self):
# for backward compatibility
return self._iter
def __enter__(self):
_CURRENT_STORAGE_STACK.append(self)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
assert _CURRENT_STORAGE_STACK[-1] == self
_CURRENT_STORAGE_STACK.pop()
@contextmanager
def name_scope(self, name):
"""
Yields:
A context within which all the events added to this storage
will be prefixed by the name scope.
"""
old_prefix = self._current_prefix
self._current_prefix = name.rstrip("/") + "/"
yield
self._current_prefix = old_prefix
def clear_images(self):
"""
Delete all the stored images for visualization. This should be called
after images are written to tensorboard.
"""
self._vis_data = []
def clear_histograms(self):
"""
Delete all the stored histograms for visualization.
This should be called after histograms are written to tensorboard.
"""
self._histograms = []
# Copyright (c) Facebook, Inc. and its affiliates.
from iopath.common.file_io import HTTPURLHandler, OneDrivePathHandler, PathHandler
from iopath.common.file_io import PathManager as PathManagerBase
__all__ = ["PathManager", "PathHandler"]
PathManager = PathManagerBase()
"""
This is a detectron2 project-specific PathManager.
We try to stay away from global PathManager in fvcore as it
introduces potential conflicts among other libraries.
"""
class Detectron2Handler(PathHandler):
"""
Resolve anything that's hosted under detectron2's namespace.
"""
PREFIX = "detectron2://"
S3_DETECTRON2_PREFIX = "https://dl.fbaipublicfiles.com/detectron2/"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path, **kwargs):
name = path[len(self.PREFIX) :]
return PathManager.get_local_path(self.S3_DETECTRON2_PREFIX + name, **kwargs)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(
self.S3_DETECTRON2_PREFIX + path[len(self.PREFIX) :], mode, **kwargs
)
PathManager.register_handler(HTTPURLHandler())
PathManager.register_handler(OneDrivePathHandler())
PathManager.register_handler(Detectron2Handler())
# Copyright (c) Facebook, Inc. and its affiliates.
import atexit
import functools
import logging
import os
import sys
import time
from collections import Counter
import torch
from tabulate import tabulate
from termcolor import colored
from detectron2.utils.file_io import PathManager
__all__ = ["setup_logger", "log_first_n", "log_every_n", "log_every_n_seconds"]
D2_LOG_BUFFER_SIZE_KEY: str = "D2_LOG_BUFFER_SIZE"
DEFAULT_LOG_BUFFER_SIZE: int = 1024 * 1024 # 1MB
class _ColorfulFormatter(logging.Formatter):
def __init__(self, *args, **kwargs):
self._root_name = kwargs.pop("root_name") + "."
self._abbrev_name = kwargs.pop("abbrev_name", "")
if len(self._abbrev_name):
self._abbrev_name = self._abbrev_name + "."
super(_ColorfulFormatter, self).__init__(*args, **kwargs)
def formatMessage(self, record):
record.name = record.name.replace(self._root_name, self._abbrev_name)
log = super(_ColorfulFormatter, self).formatMessage(record)
if record.levelno == logging.WARNING:
prefix = colored("WARNING", "red", attrs=["blink"])
elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:
prefix = colored("ERROR", "red", attrs=["blink", "underline"])
else:
return log
return prefix + " " + log
@functools.lru_cache() # so that calling setup_logger multiple times won't add many handlers
def setup_logger(
output=None,
distributed_rank=0,
*,
color=True,
name="detectron2",
abbrev_name=None,
enable_propagation: bool = False,
configure_stdout: bool = True
):
"""
Initialize the detectron2 logger and set its verbosity level to "DEBUG".
Args:
output (str): a file name or a directory to save log. If None, will not save log file.
If ends with ".txt" or ".log", assumed to be a file name.
Otherwise, logs will be saved to `output/log.txt`.
name (str): the root module name of this logger
abbrev_name (str): an abbreviation of the module, to avoid long names in logs.
Set to "" to not log the root module in logs.
By default, will abbreviate "detectron2" to "d2" and leave other
modules unchanged.
enable_propagation (bool): whether to propagate logs to the parent logger.
configure_stdout (bool): whether to configure logging to stdout.
Returns:
logging.Logger: a logger
"""
logger = logging.getLogger(name)
logger.setLevel(logging.DEBUG)
logger.propagate = enable_propagation
if abbrev_name is None:
abbrev_name = "d2" if name == "detectron2" else name
plain_formatter = logging.Formatter(
"[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%m/%d %H:%M:%S"
)
# stdout logging: master only
if configure_stdout and distributed_rank == 0:
ch = logging.StreamHandler(stream=sys.stdout)
ch.setLevel(logging.DEBUG)
if color:
formatter = _ColorfulFormatter(
colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s",
datefmt="%m/%d %H:%M:%S",
root_name=name,
abbrev_name=str(abbrev_name),
)
else:
formatter = plain_formatter
ch.setFormatter(formatter)
logger.addHandler(ch)
# file logging: all workers
if output is not None:
if output.endswith(".txt") or output.endswith(".log"):
filename = output
else:
filename = os.path.join(output, "log.txt")
if distributed_rank > 0:
filename = filename + ".rank{}".format(distributed_rank)
PathManager.mkdirs(os.path.dirname(filename))
fh = logging.StreamHandler(_cached_log_stream(filename))
fh.setLevel(logging.DEBUG)
fh.setFormatter(plain_formatter)
logger.addHandler(fh)
return logger
# cache the opened file object, so that different calls to `setup_logger`
# with the same file name can safely write to the same file.
@functools.lru_cache(maxsize=None)
def _cached_log_stream(filename):
# use 1K buffer if writing to cloud storage
io = PathManager.open(filename, "a", buffering=_get_log_stream_buffer_size(filename))
atexit.register(io.close)
return io
def _get_log_stream_buffer_size(filename: str) -> int:
if "://" not in filename:
# Local file, no extra caching is necessary
return -1
# Remote file requires a larger cache to avoid many small writes.
if D2_LOG_BUFFER_SIZE_KEY in os.environ:
return int(os.environ[D2_LOG_BUFFER_SIZE_KEY])
return DEFAULT_LOG_BUFFER_SIZE
"""
Below are some other convenient logging methods.
They are mainly adopted from
https://github.com/abseil/abseil-py/blob/master/absl/logging/__init__.py
"""
def _find_caller():
"""
Returns:
str: module name of the caller
tuple: a hashable key to be used to identify different callers
"""
frame = sys._getframe(2)
while frame:
code = frame.f_code
if os.path.join("utils", "logger.") not in code.co_filename:
mod_name = frame.f_globals["__name__"]
if mod_name == "__main__":
mod_name = "detectron2"
return mod_name, (code.co_filename, frame.f_lineno, code.co_name)
frame = frame.f_back
_LOG_COUNTER = Counter()
_LOG_TIMER = {}
def log_first_n(lvl, msg, n=1, *, name=None, key="caller"):
"""
Log only for the first n times.
Args:
lvl (int): the logging level
msg (str):
n (int):
name (str): name of the logger to use. Will use the caller's module by default.
key (str or tuple[str]): the string(s) can be one of "caller" or
"message", which defines how to identify duplicated logs.
For example, if called with `n=1, key="caller"`, this function
will only log the first call from the same caller, regardless of
the message content.
If called with `n=1, key="message"`, this function will log the
same content only once, even if they are called from different places.
If called with `n=1, key=("caller", "message")`, this function
will not log only if the same caller has logged the same message before.
"""
if isinstance(key, str):
key = (key,)
assert len(key) > 0
caller_module, caller_key = _find_caller()
hash_key = ()
if "caller" in key:
hash_key = hash_key + caller_key
if "message" in key:
hash_key = hash_key + (msg,)
_LOG_COUNTER[hash_key] += 1
if _LOG_COUNTER[hash_key] <= n:
logging.getLogger(name or caller_module).log(lvl, msg)
def log_every_n(lvl, msg, n=1, *, name=None):
"""
Log once per n times.
Args:
lvl (int): the logging level
msg (str):
n (int):
name (str): name of the logger to use. Will use the caller's module by default.
"""
caller_module, key = _find_caller()
_LOG_COUNTER[key] += 1
if n == 1 or _LOG_COUNTER[key] % n == 1:
logging.getLogger(name or caller_module).log(lvl, msg)
def log_every_n_seconds(lvl, msg, n=1, *, name=None):
"""
Log no more than once per n seconds.
Args:
lvl (int): the logging level
msg (str):
n (int):
name (str): name of the logger to use. Will use the caller's module by default.
"""
caller_module, key = _find_caller()
last_logged = _LOG_TIMER.get(key, None)
current_time = time.time()
if last_logged is None or current_time - last_logged >= n:
logging.getLogger(name or caller_module).log(lvl, msg)
_LOG_TIMER[key] = current_time
def create_small_table(small_dict):
"""
Create a small table using the keys of small_dict as headers. This is only
suitable for small dictionaries.
Args:
small_dict (dict): a result dictionary of only a few items.
Returns:
str: the table as a string.
"""
keys, values = tuple(zip(*small_dict.items()))
table = tabulate(
[values],
headers=keys,
tablefmt="pipe",
floatfmt=".3f",
stralign="center",
numalign="center",
)
return table
def _log_api_usage(identifier: str):
"""
Internal function used to log the usage of different detectron2 components
inside facebook's infra.
"""
torch._C._log_api_usage_once("detectron2." + identifier)
# Copyright (c) Facebook, Inc. and its affiliates.
import logging
from contextlib import contextmanager
from functools import wraps
import torch
__all__ = ["retry_if_cuda_oom"]
@contextmanager
def _ignore_torch_cuda_oom():
"""
A context which ignores CUDA OOM exception from pytorch.
"""
try:
yield
except RuntimeError as e:
# NOTE: the string may change?
if "CUDA out of memory. " in str(e):
pass
else:
raise
def retry_if_cuda_oom(func):
"""
Makes a function retry itself after encountering
pytorch's CUDA OOM error.
It will first retry after calling `torch.cuda.empty_cache()`.
If that still fails, it will then retry by trying to convert inputs to CPUs.
In this case, it expects the function to dispatch to CPU implementation.
The return values may become CPU tensors as well and it's user's
responsibility to convert it back to CUDA tensor if needed.
Args:
func: a stateless callable that takes tensor-like objects as arguments
Returns:
a callable which retries `func` if OOM is encountered.
Examples:
::
output = retry_if_cuda_oom(some_torch_function)(input1, input2)
# output may be on CPU even if inputs are on GPU
Note:
1. When converting inputs to CPU, it will only look at each argument and check
if it has `.device` and `.to` for conversion. Nested structures of tensors
are not supported.
2. Since the function might be called more than once, it has to be
stateless.
"""
def maybe_to_cpu(x):
try:
like_gpu_tensor = x.device.type == "cuda" and hasattr(x, "to")
except AttributeError:
like_gpu_tensor = False
if like_gpu_tensor:
return x.to(device="cpu")
else:
return x
@wraps(func)
def wrapped(*args, **kwargs):
with _ignore_torch_cuda_oom():
return func(*args, **kwargs)
# Clear cache and retry
torch.cuda.empty_cache()
with _ignore_torch_cuda_oom():
return func(*args, **kwargs)
# Try on CPU. This slows down the code significantly, therefore print a notice.
logger = logging.getLogger(__name__)
logger.info("Attempting to copy inputs of {} to CPU due to CUDA OOM".format(str(func)))
new_args = (maybe_to_cpu(x) for x in args)
new_kwargs = {k: maybe_to_cpu(v) for k, v in kwargs.items()}
return func(*new_args, **new_kwargs)
return wrapped
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Any
import pydoc
from fvcore.common.registry import Registry # for backward compatibility.
"""
``Registry`` and `locate` provide ways to map a string (typically found
in config files) to callable objects.
"""
__all__ = ["Registry", "locate"]
def _convert_target_to_string(t: Any) -> str:
"""
Inverse of ``locate()``.
Args:
t: any object with ``__module__`` and ``__qualname__``
"""
module, qualname = t.__module__, t.__qualname__
# Compress the path to this object, e.g. ``module.submodule._impl.class``
# may become ``module.submodule.class``, if the later also resolves to the same
# object. This simplifies the string, and also is less affected by moving the
# class implementation.
module_parts = module.split(".")
for k in range(1, len(module_parts)):
prefix = ".".join(module_parts[:k])
candidate = f"{prefix}.{qualname}"
try:
if locate(candidate) is t:
return candidate
except ImportError:
pass
return f"{module}.{qualname}"
def locate(name: str) -> Any:
"""
Locate and return an object ``x`` using an input string ``{x.__module__}.{x.__qualname__}``,
such as "module.submodule.class_name".
Raise Exception if it cannot be found.
"""
obj = pydoc.locate(name)
# Some cases (e.g. torch.optim.sgd.SGD) not handled correctly
# by pydoc.locate. Try a private function from hydra.
if obj is None:
try:
# from hydra.utils import get_method - will print many errors
from hydra.utils import _locate
except ImportError as e:
raise ImportError(f"Cannot dynamically locate object {name}!") from e
else:
obj = _locate(name) # it raises if fails
return obj
# Copyright (c) Facebook, Inc. and its affiliates.
import cloudpickle
class PicklableWrapper:
"""
Wrap an object to make it more picklable, note that it uses
heavy weight serialization libraries that are slower than pickle.
It's best to use it only on closures (which are usually not picklable).
This is a simplified version of
https://github.com/joblib/joblib/blob/master/joblib/externals/loky/cloudpickle_wrapper.py
"""
def __init__(self, obj):
while isinstance(obj, PicklableWrapper):
# Wrapping an object twice is no-op
obj = obj._obj
self._obj = obj
def __reduce__(self):
s = cloudpickle.dumps(self._obj)
return cloudpickle.loads, (s,)
def __call__(self, *args, **kwargs):
return self._obj(*args, **kwargs)
def __getattr__(self, attr):
# Ensure that the wrapped object can be used seamlessly as the previous object.
if attr not in ["_obj"]:
return getattr(self._obj, attr)
return getattr(self, attr)
# Copyright (c) Facebook, Inc. and its affiliates.
import io
import numpy as np
import os
import re
import tempfile
import unittest
from typing import Callable
import torch
import torch.onnx.symbolic_helper as sym_help
from torch._C import ListType
from torch.onnx import register_custom_op_symbolic
from detectron2 import model_zoo
from detectron2.config import CfgNode, LazyConfig, instantiate
from detectron2.data import DatasetCatalog
from detectron2.data.detection_utils import read_image
from detectron2.modeling import build_model
from detectron2.structures import Boxes, Instances, ROIMasks
from detectron2.utils.file_io import PathManager
from detectron2.utils.torch_version_utils import min_torch_version
"""
Internal utilities for tests. Don't use except for writing tests.
"""
def get_model_no_weights(config_path):
"""
Like model_zoo.get, but do not load any weights (even pretrained)
"""
cfg = model_zoo.get_config(config_path)
if isinstance(cfg, CfgNode):
if not torch.cuda.is_available():
cfg.MODEL.DEVICE = "cpu"
return build_model(cfg)
else:
return instantiate(cfg.model)
def random_boxes(num_boxes, max_coord=100, device="cpu"):
"""
Create a random Nx4 boxes tensor, with coordinates < max_coord.
"""
boxes = torch.rand(num_boxes, 4, device=device) * (max_coord * 0.5)
boxes.clamp_(min=1.0) # tiny boxes cause numerical instability in box regression
# Note: the implementation of this function in torchvision is:
# boxes[:, 2:] += torch.rand(N, 2) * 100
# but it does not guarantee non-negative widths/heights constraints:
# boxes[:, 2] >= boxes[:, 0] and boxes[:, 3] >= boxes[:, 1]:
boxes[:, 2:] += boxes[:, :2]
return boxes
def get_sample_coco_image(tensor=True):
"""
Args:
tensor (bool): if True, returns 3xHxW tensor.
else, returns a HxWx3 numpy array.
Returns:
an image, in BGR color.
"""
try:
file_name = DatasetCatalog.get("coco_2017_val_100")[0]["file_name"]
if not PathManager.exists(file_name):
raise FileNotFoundError()
except IOError:
# for public CI to run
file_name = PathManager.get_local_path(
"http://images.cocodataset.org/train2017/000000000009.jpg"
)
ret = read_image(file_name, format="BGR")
if tensor:
ret = torch.from_numpy(np.ascontiguousarray(ret.transpose(2, 0, 1)))
return ret
def convert_scripted_instances(instances):
"""
Convert a scripted Instances object to a regular :class:`Instances` object
"""
assert hasattr(
instances, "image_size"
), f"Expect an Instances object, but got {type(instances)}!"
ret = Instances(instances.image_size)
for name in instances._field_names:
val = getattr(instances, "_" + name, None)
if val is not None:
ret.set(name, val)
return ret
def assert_instances_allclose(input, other, *, rtol=1e-5, msg="", size_as_tensor=False):
"""
Args:
input, other (Instances):
size_as_tensor: compare image_size of the Instances as tensors (instead of tuples).
Useful for comparing outputs of tracing.
"""
if not isinstance(input, Instances):
input = convert_scripted_instances(input)
if not isinstance(other, Instances):
other = convert_scripted_instances(other)
if not msg:
msg = "Two Instances are different! "
else:
msg = msg.rstrip() + " "
size_error_msg = msg + f"image_size is {input.image_size} vs. {other.image_size}!"
if size_as_tensor:
assert torch.equal(
torch.tensor(input.image_size), torch.tensor(other.image_size)
), size_error_msg
else:
assert input.image_size == other.image_size, size_error_msg
fields = sorted(input.get_fields().keys())
fields_other = sorted(other.get_fields().keys())
assert fields == fields_other, msg + f"Fields are {fields} vs {fields_other}!"
for f in fields:
val1, val2 = input.get(f), other.get(f)
if isinstance(val1, (Boxes, ROIMasks)):
# boxes in the range of O(100) and can have a larger tolerance
assert torch.allclose(val1.tensor, val2.tensor, atol=100 * rtol), (
msg + f"Field {f} differs too much!"
)
elif isinstance(val1, torch.Tensor):
if val1.dtype.is_floating_point:
mag = torch.abs(val1).max().cpu().item()
assert torch.allclose(val1, val2, atol=mag * rtol), (
msg + f"Field {f} differs too much!"
)
else:
assert torch.equal(val1, val2), msg + f"Field {f} is different!"
else:
raise ValueError(f"Don't know how to compare type {type(val1)}")
def reload_script_model(module):
"""
Save a jit module and load it back.
Similar to the `getExportImportCopy` function in torch/testing/
"""
buffer = io.BytesIO()
torch.jit.save(module, buffer)
buffer.seek(0)
return torch.jit.load(buffer)
def reload_lazy_config(cfg):
"""
Save an object by LazyConfig.save and load it back.
This is used to test that a config still works the same after
serialization/deserialization.
"""
with tempfile.TemporaryDirectory(prefix="detectron2") as d:
fname = os.path.join(d, "d2_cfg_test.yaml")
LazyConfig.save(cfg, fname)
return LazyConfig.load(fname)
def has_dynamic_axes(onnx_model):
"""
Return True when all ONNX input/output have only dynamic axes for all ranks
"""
return all(
not dim.dim_param.isnumeric()
for inp in onnx_model.graph.input
for dim in inp.type.tensor_type.shape.dim
) and all(
not dim.dim_param.isnumeric()
for out in onnx_model.graph.output
for dim in out.type.tensor_type.shape.dim
)
def register_custom_op_onnx_export(
opname: str, symbolic_fn: Callable, opset_version: int, min_version: str
) -> None:
"""
Register `symbolic_fn` as PyTorch's symbolic `opname`-`opset_version` for ONNX export.
The registration is performed only when current PyTorch's version is < `min_version.`
IMPORTANT: symbolic must be manually unregistered after the caller function returns
"""
if min_torch_version(min_version):
return
register_custom_op_symbolic(opname, symbolic_fn, opset_version)
print(f"_register_custom_op_onnx_export({opname}, {opset_version}) succeeded.")
def unregister_custom_op_onnx_export(opname: str, opset_version: int, min_version: str) -> None:
"""
Unregister PyTorch's symbolic `opname`-`opset_version` for ONNX export.
The un-registration is performed only when PyTorch's version is < `min_version`
IMPORTANT: The symbolic must have been manually registered by the caller, otherwise
the incorrect symbolic may be unregistered instead.
"""
# TODO: _unregister_custom_op_symbolic is introduced PyTorch>=1.10
# Remove after PyTorch 1.10+ is used by ALL detectron2's CI
try:
from torch.onnx import unregister_custom_op_symbolic as _unregister_custom_op_symbolic
except ImportError:
def _unregister_custom_op_symbolic(symbolic_name, opset_version):
import torch.onnx.symbolic_registry as sym_registry
from torch.onnx.symbolic_helper import _onnx_main_opset, _onnx_stable_opsets
def _get_ns_op_name_from_custom_op(symbolic_name):
try:
from torch.onnx.utils import get_ns_op_name_from_custom_op
ns, op_name = get_ns_op_name_from_custom_op(symbolic_name)
except ImportError as import_error:
if not bool(
re.match(r"^[a-zA-Z0-9-_]*::[a-zA-Z-_]+[a-zA-Z0-9-_]*$", symbolic_name)
):
raise ValueError(
f"Invalid symbolic name {symbolic_name}. Must be `domain::name`"
) from import_error
ns, op_name = symbolic_name.split("::")
if ns == "onnx":
raise ValueError(f"{ns} domain cannot be modified.") from import_error
if ns == "aten":
ns = ""
return ns, op_name
def _unregister_op(opname: str, domain: str, version: int):
try:
sym_registry.unregister_op(op_name, ns, ver)
except AttributeError as attribute_error:
if sym_registry.is_registered_op(opname, domain, version):
del sym_registry._registry[(domain, version)][opname]
if not sym_registry._registry[(domain, version)]:
del sym_registry._registry[(domain, version)]
else:
raise RuntimeError(
f"The opname {opname} is not registered."
) from attribute_error
ns, op_name = _get_ns_op_name_from_custom_op(symbolic_name)
for ver in _onnx_stable_opsets + [_onnx_main_opset]:
if ver >= opset_version:
_unregister_op(op_name, ns, ver)
if min_torch_version(min_version):
return
_unregister_custom_op_symbolic(opname, opset_version)
print(f"_unregister_custom_op_onnx_export({opname}, {opset_version}) succeeded.")
skipIfOnCPUCI = unittest.skipIf(
os.environ.get("CI") and not torch.cuda.is_available(),
"The test is too slow on CPUs and will be executed on CircleCI's GPU jobs.",
)
def skipIfUnsupportedMinOpsetVersion(min_opset_version, current_opset_version=None):
"""
Skips tests for ONNX Opset versions older than min_opset_version.
"""
def skip_dec(func):
def wrapper(self):
try:
opset_version = self.opset_version
except AttributeError:
opset_version = current_opset_version
if opset_version < min_opset_version:
raise unittest.SkipTest(
f"Unsupported opset_version {opset_version}"
f", required is {min_opset_version}"
)
return func(self)
return wrapper
return skip_dec
def skipIfUnsupportedMinTorchVersion(min_version):
"""
Skips tests for PyTorch versions older than min_version.
"""
reason = f"module 'torch' has __version__ {torch.__version__}" f", required is: {min_version}"
return unittest.skipIf(not min_torch_version(min_version), reason)
# TODO: Remove after PyTorch 1.11.1+ is used by detectron2's CI
def _pytorch1111_symbolic_opset9_to(g, self, *args):
"""aten::to() symbolic that must be used for testing with PyTorch < 1.11.1."""
def is_aten_to_device_only(args):
if len(args) == 4:
# aten::to(Tensor, Device, bool, bool, memory_format)
return (
args[0].node().kind() == "prim::device"
or args[0].type().isSubtypeOf(ListType.ofInts())
or (
sym_help._is_value(args[0])
and args[0].node().kind() == "onnx::Constant"
and isinstance(args[0].node()["value"], str)
)
)
elif len(args) == 5:
# aten::to(Tensor, Device, ScalarType, bool, bool, memory_format)
# When dtype is None, this is a aten::to(device) call
dtype = sym_help._get_const(args[1], "i", "dtype")
return dtype is None
elif len(args) in (6, 7):
# aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format)
# aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format)
# When dtype is None, this is a aten::to(device) call
dtype = sym_help._get_const(args[0], "i", "dtype")
return dtype is None
return False
# ONNX doesn't have a concept of a device, so we ignore device-only casts
if is_aten_to_device_only(args):
return self
if len(args) == 4:
# TestONNXRuntime::test_ones_bool shows args[0] of aten::to can be onnx::Constant[Tensor]
# In this case, the constant value is a tensor not int,
# so sym_help._maybe_get_const(args[0], 'i') would not work.
dtype = args[0]
if sym_help._is_value(args[0]) and args[0].node().kind() == "onnx::Constant":
tval = args[0].node()["value"]
if isinstance(tval, torch.Tensor):
if len(tval.shape) == 0:
tval = tval.item()
dtype = int(tval)
else:
dtype = tval
if sym_help._is_value(dtype) or isinstance(dtype, torch.Tensor):
# aten::to(Tensor, Tensor, bool, bool, memory_format)
dtype = args[0].type().scalarType()
return g.op("Cast", self, to_i=sym_help.cast_pytorch_to_onnx[dtype])
else:
# aten::to(Tensor, ScalarType, bool, bool, memory_format)
# memory_format is ignored
return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])
elif len(args) == 5:
# aten::to(Tensor, Device, ScalarType, bool, bool, memory_format)
dtype = sym_help._get_const(args[1], "i", "dtype")
# memory_format is ignored
return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])
elif len(args) == 6:
# aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format)
dtype = sym_help._get_const(args[0], "i", "dtype")
# Layout, device and memory_format are ignored
return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])
elif len(args) == 7:
# aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format)
dtype = sym_help._get_const(args[0], "i", "dtype")
# Layout, device and memory_format are ignored
return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])
else:
return sym_help._onnx_unsupported("Unknown aten::to signature")
# TODO: Remove after PyTorch 1.11.1+ is used by detectron2's CI
def _pytorch1111_symbolic_opset9_repeat_interleave(g, self, repeats, dim=None, output_size=None):
# from torch.onnx.symbolic_helper import ScalarType
from torch.onnx.symbolic_opset9 import expand, unsqueeze
input = self
# if dim is None flatten
# By default, use the flattened input array, and return a flat output array
if sym_help._is_none(dim):
input = sym_help._reshape_helper(g, self, g.op("Constant", value_t=torch.tensor([-1])))
dim = 0
else:
dim = sym_help._maybe_get_scalar(dim)
repeats_dim = sym_help._get_tensor_rank(repeats)
repeats_sizes = sym_help._get_tensor_sizes(repeats)
input_sizes = sym_help._get_tensor_sizes(input)
if repeats_dim is None:
raise RuntimeError(
"Unsupported: ONNX export of repeat_interleave for unknown " "repeats rank."
)
if repeats_sizes is None:
raise RuntimeError(
"Unsupported: ONNX export of repeat_interleave for unknown " "repeats size."
)
if input_sizes is None:
raise RuntimeError(
"Unsupported: ONNX export of repeat_interleave for unknown " "input size."
)
input_sizes_temp = input_sizes.copy()
for idx, input_size in enumerate(input_sizes):
if input_size is None:
input_sizes[idx], input_sizes_temp[idx] = 0, -1
# Cases where repeats is an int or single value tensor
if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1):
if not sym_help._is_tensor(repeats):
repeats = g.op("Constant", value_t=torch.LongTensor(repeats))
if input_sizes[dim] == 0:
return sym_help._onnx_opset_unsupported_detailed(
"repeat_interleave",
9,
13,
"Unsupported along dimension with unknown input size",
)
else:
reps = input_sizes[dim]
repeats = expand(g, repeats, g.op("Constant", value_t=torch.tensor([reps])), None)
# Cases where repeats is a 1 dim Tensor
elif repeats_dim == 1:
if input_sizes[dim] == 0:
return sym_help._onnx_opset_unsupported_detailed(
"repeat_interleave",
9,
13,
"Unsupported along dimension with unknown input size",
)
if repeats_sizes[0] is None:
return sym_help._onnx_opset_unsupported_detailed(
"repeat_interleave", 9, 13, "Unsupported for cases with dynamic repeats"
)
assert (
repeats_sizes[0] == input_sizes[dim]
), "repeats must have the same size as input along dim"
reps = repeats_sizes[0]
else:
raise RuntimeError("repeats must be 0-dim or 1-dim tensor")
final_splits = list()
r_splits = sym_help._repeat_interleave_split_helper(g, repeats, reps, 0)
if isinstance(r_splits, torch._C.Value):
r_splits = [r_splits]
i_splits = sym_help._repeat_interleave_split_helper(g, input, reps, dim)
if isinstance(i_splits, torch._C.Value):
i_splits = [i_splits]
input_sizes[dim], input_sizes_temp[dim] = -1, 1
for idx, r_split in enumerate(r_splits):
i_split = unsqueeze(g, i_splits[idx], dim + 1)
r_concat = [
g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[: dim + 1])),
r_split,
g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[dim + 1 :])),
]
r_concat = g.op("Concat", *r_concat, axis_i=0)
i_split = expand(g, i_split, r_concat, None)
i_split = sym_help._reshape_helper(
g,
i_split,
g.op("Constant", value_t=torch.LongTensor(input_sizes)),
allowzero=0,
)
final_splits.append(i_split)
return g.op("Concat", *final_splits, axis_i=dim)
from packaging import version
def min_torch_version(min_version: str) -> bool:
"""
Returns True when torch's version is at least `min_version`.
"""
try:
import torch
except ImportError:
return False
installed_version = version.parse(torch.__version__.split("+")[0])
min_version = version.parse(min_version)
return installed_version >= min_version
import inspect
import torch
from detectron2.utils.env import TORCH_VERSION
try:
from torch.fx._symbolic_trace import is_fx_tracing as is_fx_tracing_current
tracing_current_exists = True
except ImportError:
tracing_current_exists = False
try:
from torch.fx._symbolic_trace import _orig_module_call
tracing_legacy_exists = True
except ImportError:
tracing_legacy_exists = False
@torch.jit.ignore
def is_fx_tracing_legacy() -> bool:
"""
Returns a bool indicating whether torch.fx is currently symbolically tracing a module.
Can be useful for gating module logic that is incompatible with symbolic tracing.
"""
return torch.nn.Module.__call__ is not _orig_module_call
def is_fx_tracing() -> bool:
"""Returns whether execution is currently in
Torch FX tracing mode"""
if torch.jit.is_scripting():
return False
if TORCH_VERSION >= (1, 10) and tracing_current_exists:
return is_fx_tracing_current()
elif tracing_legacy_exists:
return is_fx_tracing_legacy()
else:
# Can't find either current or legacy tracing indication code.
# Enabling this assert_fx_safe() call regardless of tracing status.
return False
def assert_fx_safe(condition: bool, message: str) -> torch.Tensor:
"""An FX-tracing safe version of assert.
Avoids erroneous type assertion triggering when types are masked inside
an fx.proxy.Proxy object during tracing.
Args: condition - either a boolean expression or a string representing
the condition to test. If this assert triggers an exception when tracing
due to dynamic control flow, try encasing the expression in quotation
marks and supplying it as a string."""
# Must return a concrete tensor for compatibility with PyTorch <=1.8.
# If <=1.8 compatibility is not needed, return type can be converted to None
if torch.jit.is_scripting() or is_fx_tracing():
return torch.zeros(1)
return _do_assert_fx_safe(condition, message)
def _do_assert_fx_safe(condition: bool, message: str) -> torch.Tensor:
try:
if isinstance(condition, str):
caller_frame = inspect.currentframe().f_back
torch._assert(eval(condition, caller_frame.f_globals, caller_frame.f_locals), message)
return torch.ones(1)
else:
torch._assert(condition, message)
return torch.ones(1)
except torch.fx.proxy.TraceError as e:
print(
"Found a non-FX compatible assertion. Skipping the check. Failure is shown below"
+ str(e)
)
# Copyright (c) Facebook, Inc. and its affiliates.
import numpy as np
from typing import List
import pycocotools.mask as mask_util
from detectron2.structures import Instances
from detectron2.utils.visualizer import (
ColorMode,
Visualizer,
_create_text_labels,
_PanopticPrediction,
)
from .colormap import random_color, random_colors
class _DetectedInstance:
"""
Used to store data about detected objects in video frame,
in order to transfer color to objects in the future frames.
Attributes:
label (int):
bbox (tuple[float]):
mask_rle (dict):
color (tuple[float]): RGB colors in range (0, 1)
ttl (int): time-to-live for the instance. For example, if ttl=2,
the instance color can be transferred to objects in the next two frames.
"""
__slots__ = ["label", "bbox", "mask_rle", "color", "ttl"]
def __init__(self, label, bbox, mask_rle, color, ttl):
self.label = label
self.bbox = bbox
self.mask_rle = mask_rle
self.color = color
self.ttl = ttl
class VideoVisualizer:
def __init__(self, metadata, instance_mode=ColorMode.IMAGE):
"""
Args:
metadata (MetadataCatalog): image metadata.
"""
self.metadata = metadata
self._old_instances = []
assert instance_mode in [
ColorMode.IMAGE,
ColorMode.IMAGE_BW,
], "Other mode not supported yet."
self._instance_mode = instance_mode
self._max_num_instances = self.metadata.get("max_num_instances", 74)
self._assigned_colors = {}
self._color_pool = random_colors(self._max_num_instances, rgb=True, maximum=1)
self._color_idx_set = set(range(len(self._color_pool)))
def draw_instance_predictions(self, frame, predictions):
"""
Draw instance-level prediction results on an image.
Args:
frame (ndarray): an RGB image of shape (H, W, C), in the range [0, 255].
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
Returns:
output (VisImage): image object with visualizations.
"""
frame_visualizer = Visualizer(frame, self.metadata)
num_instances = len(predictions)
if num_instances == 0:
return frame_visualizer.output
boxes = predictions.pred_boxes.tensor.numpy() if predictions.has("pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes.numpy() if predictions.has("pred_classes") else None
keypoints = predictions.pred_keypoints if predictions.has("pred_keypoints") else None
colors = predictions.COLOR if predictions.has("COLOR") else [None] * len(predictions)
periods = predictions.ID_period if predictions.has("ID_period") else None
period_threshold = self.metadata.get("period_threshold", 0)
visibilities = (
[True] * len(predictions)
if periods is None
else [x > period_threshold for x in periods]
)
if predictions.has("pred_masks"):
masks = predictions.pred_masks
# mask IOU is not yet enabled
# masks_rles = mask_util.encode(np.asarray(masks.permute(1, 2, 0), order="F"))
# assert len(masks_rles) == num_instances
else:
masks = None
if not predictions.has("COLOR"):
if predictions.has("ID"):
colors = self._assign_colors_by_id(predictions)
else:
# ToDo: clean old assign color method and use a default tracker to assign id
detected = [
_DetectedInstance(classes[i], boxes[i], mask_rle=None, color=colors[i], ttl=8)
for i in range(num_instances)
]
colors = self._assign_colors(detected)
labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None))
if self._instance_mode == ColorMode.IMAGE_BW:
# any() returns uint8 tensor
frame_visualizer.output.reset_image(
frame_visualizer._create_grayscale_image(
(masks.any(dim=0) > 0).numpy() if masks is not None else None
)
)
alpha = 0.3
else:
alpha = 0.5
labels = (
None
if labels is None
else [y[0] for y in filter(lambda x: x[1], zip(labels, visibilities))]
) # noqa
assigned_colors = (
None
if colors is None
else [y[0] for y in filter(lambda x: x[1], zip(colors, visibilities))]
) # noqa
frame_visualizer.overlay_instances(
boxes=None if masks is not None else boxes[visibilities], # boxes are a bit distracting
masks=None if masks is None else masks[visibilities],
labels=labels,
keypoints=None if keypoints is None else keypoints[visibilities],
assigned_colors=assigned_colors,
alpha=alpha,
)
return frame_visualizer.output
def draw_sem_seg(self, frame, sem_seg, area_threshold=None):
"""
Args:
sem_seg (ndarray or Tensor): semantic segmentation of shape (H, W),
each value is the integer label.
area_threshold (Optional[int]): only draw segmentations larger than the threshold
"""
# don't need to do anything special
frame_visualizer = Visualizer(frame, self.metadata)
frame_visualizer.draw_sem_seg(sem_seg, area_threshold=None)
return frame_visualizer.output
def draw_panoptic_seg_predictions(
self, frame, panoptic_seg, segments_info, area_threshold=None, alpha=0.5
):
frame_visualizer = Visualizer(frame, self.metadata)
pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata)
if self._instance_mode == ColorMode.IMAGE_BW:
frame_visualizer.output.reset_image(
frame_visualizer._create_grayscale_image(pred.non_empty_mask())
)
# draw mask for all semantic segments first i.e. "stuff"
for mask, sinfo in pred.semantic_masks():
category_idx = sinfo["category_id"]
try:
mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]]
except AttributeError:
mask_color = None
frame_visualizer.draw_binary_mask(
mask,
color=mask_color,
text=self.metadata.stuff_classes[category_idx],
alpha=alpha,
area_threshold=area_threshold,
)
all_instances = list(pred.instance_masks())
if len(all_instances) == 0:
return frame_visualizer.output
# draw mask for all instances second
masks, sinfo = list(zip(*all_instances))
num_instances = len(masks)
masks_rles = mask_util.encode(
np.asarray(np.asarray(masks).transpose(1, 2, 0), dtype=np.uint8, order="F")
)
assert len(masks_rles) == num_instances
category_ids = [x["category_id"] for x in sinfo]
detected = [
_DetectedInstance(category_ids[i], bbox=None, mask_rle=masks_rles[i], color=None, ttl=8)
for i in range(num_instances)
]
colors = self._assign_colors(detected)
labels = [self.metadata.thing_classes[k] for k in category_ids]
frame_visualizer.overlay_instances(
boxes=None,
masks=masks,
labels=labels,
keypoints=None,
assigned_colors=colors,
alpha=alpha,
)
return frame_visualizer.output
def _assign_colors(self, instances):
"""
Naive tracking heuristics to assign same color to the same instance,
will update the internal state of tracked instances.
Returns:
list[tuple[float]]: list of colors.
"""
# Compute iou with either boxes or masks:
is_crowd = np.zeros((len(instances),), dtype=bool)
if instances[0].bbox is None:
assert instances[0].mask_rle is not None
# use mask iou only when box iou is None
# because box seems good enough
rles_old = [x.mask_rle for x in self._old_instances]
rles_new = [x.mask_rle for x in instances]
ious = mask_util.iou(rles_old, rles_new, is_crowd)
threshold = 0.5
else:
boxes_old = [x.bbox for x in self._old_instances]
boxes_new = [x.bbox for x in instances]
ious = mask_util.iou(boxes_old, boxes_new, is_crowd)
threshold = 0.6
if len(ious) == 0:
ious = np.zeros((len(self._old_instances), len(instances)), dtype="float32")
# Only allow matching instances of the same label:
for old_idx, old in enumerate(self._old_instances):
for new_idx, new in enumerate(instances):
if old.label != new.label:
ious[old_idx, new_idx] = 0
matched_new_per_old = np.asarray(ious).argmax(axis=1)
max_iou_per_old = np.asarray(ious).max(axis=1)
# Try to find match for each old instance:
extra_instances = []
for idx, inst in enumerate(self._old_instances):
if max_iou_per_old[idx] > threshold:
newidx = matched_new_per_old[idx]
if instances[newidx].color is None:
instances[newidx].color = inst.color
continue
# If an old instance does not match any new instances,
# keep it for the next frame in case it is just missed by the detector
inst.ttl -= 1
if inst.ttl > 0:
extra_instances.append(inst)
# Assign random color to newly-detected instances:
for inst in instances:
if inst.color is None:
inst.color = random_color(rgb=True, maximum=1)
self._old_instances = instances[:] + extra_instances
return [d.color for d in instances]
def _assign_colors_by_id(self, instances: Instances) -> List:
colors = []
untracked_ids = set(self._assigned_colors.keys())
for id in instances.ID:
if id in self._assigned_colors:
colors.append(self._color_pool[self._assigned_colors[id]])
untracked_ids.remove(id)
else:
assert (
len(self._color_idx_set) >= 1
), f"Number of id exceeded maximum, \
max = {self._max_num_instances}"
idx = self._color_idx_set.pop()
color = self._color_pool[idx]
self._assigned_colors[id] = idx
colors.append(color)
for id in untracked_ids:
self._color_idx_set.add(self._assigned_colors[id])
del self._assigned_colors[id]
return colors
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment