File size: 5,988 Bytes
404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af |
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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 |
import os
import contextlib
import joblib
from typing import Union
from loguru import _Logger, logger
from itertools import chain
import torch
from yacs.config import CfgNode as CN
from pytorch_lightning.utilities import rank_zero_only
import cv2
import numpy as np
def lower_config(yacs_cfg):
if not isinstance(yacs_cfg, CN):
return yacs_cfg
return {k.lower(): lower_config(v) for k, v in yacs_cfg.items()}
def upper_config(dict_cfg):
if not isinstance(dict_cfg, dict):
return dict_cfg
return {k.upper(): upper_config(v) for k, v in dict_cfg.items()}
def log_on(condition, message, level):
if condition:
assert level in ["INFO", "DEBUG", "WARNING", "ERROR", "CRITICAL"]
logger.log(level, message)
def get_rank_zero_only_logger(logger: _Logger):
if rank_zero_only.rank == 0:
return logger
else:
for _level in logger._core.levels.keys():
level = _level.lower()
setattr(logger, level, lambda x: None)
logger._log = lambda x: None
return logger
def setup_gpus(gpus: Union[str, int]) -> int:
"""A temporary fix for pytorch-lighting 1.3.x"""
gpus = str(gpus)
gpu_ids = []
if "," not in gpus:
n_gpus = int(gpus)
return n_gpus if n_gpus != -1 else torch.cuda.device_count()
else:
gpu_ids = [i.strip() for i in gpus.split(",") if i != ""]
# setup environment variables
visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
if visible_devices is None:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_ids)
visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
logger.warning(
f"[Temporary Fix] manually set CUDA_VISIBLE_DEVICES when specifying gpus to use: {visible_devices}"
)
else:
logger.warning(
"[Temporary Fix] CUDA_VISIBLE_DEVICES already set by user or the main process."
)
return len(gpu_ids)
def flattenList(x):
return list(chain(*x))
@contextlib.contextmanager
def tqdm_joblib(tqdm_object):
"""Context manager to patch joblib to report into tqdm progress bar given as argument
Usage:
with tqdm_joblib(tqdm(desc="My calculation", total=10)) as progress_bar:
Parallel(n_jobs=16)(delayed(sqrt)(i**2) for i in range(10))
When iterating over a generator, directly use of tqdm is also a solutin (but monitor the task queuing, instead of finishing)
ret_vals = Parallel(n_jobs=args.world_size)(
delayed(lambda x: _compute_cov_score(pid, *x))(param)
for param in tqdm(combinations(image_ids, 2),
desc=f'Computing cov_score of [{pid}]',
total=len(image_ids)*(len(image_ids)-1)/2))
Src: https://stackoverflow.com/a/58936697
"""
class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __call__(self, *args, **kwargs):
tqdm_object.update(n=self.batch_size)
return super().__call__(*args, **kwargs)
old_batch_callback = joblib.parallel.BatchCompletionCallBack
joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
try:
yield tqdm_object
finally:
joblib.parallel.BatchCompletionCallBack = old_batch_callback
tqdm_object.close()
def draw_points(img, points, color=(0, 255, 0), radius=3):
dp = [(int(points[i, 0]), int(points[i, 1])) for i in range(points.shape[0])]
for i in range(points.shape[0]):
cv2.circle(img, dp[i], radius=radius, color=color)
return img
def draw_match(
img1,
img2,
corr1,
corr2,
inlier=[True],
color=None,
radius1=1,
radius2=1,
resize=None,
):
if resize is not None:
scale1, scale2 = [img1.shape[1] / resize[0], img1.shape[0] / resize[1]], [
img2.shape[1] / resize[0],
img2.shape[0] / resize[1],
]
img1, img2 = cv2.resize(img1, resize, interpolation=cv2.INTER_AREA), cv2.resize(
img2, resize, interpolation=cv2.INTER_AREA
)
corr1, corr2 = (
corr1 / np.asarray(scale1)[np.newaxis],
corr2 / np.asarray(scale2)[np.newaxis],
)
corr1_key = [
cv2.KeyPoint(corr1[i, 0], corr1[i, 1], radius1) for i in range(corr1.shape[0])
]
corr2_key = [
cv2.KeyPoint(corr2[i, 0], corr2[i, 1], radius2) for i in range(corr2.shape[0])
]
assert len(corr1) == len(corr2)
draw_matches = [cv2.DMatch(i, i, 0) for i in range(len(corr1))]
if color is None:
color = [(0, 255, 0) if cur_inlier else (0, 0, 255) for cur_inlier in inlier]
if len(color) == 1:
display = cv2.drawMatches(
img1,
corr1_key,
img2,
corr2_key,
draw_matches,
None,
matchColor=color[0],
singlePointColor=color[0],
flags=4,
)
else:
height, width = max(img1.shape[0], img2.shape[0]), img1.shape[1] + img2.shape[1]
display = np.zeros([height, width, 3], np.uint8)
display[: img1.shape[0], : img1.shape[1]] = img1
display[: img2.shape[0], img1.shape[1] :] = img2
for i in range(len(corr1)):
left_x, left_y, right_x, right_y = (
int(corr1[i][0]),
int(corr1[i][1]),
int(corr2[i][0] + img1.shape[1]),
int(corr2[i][1]),
)
cur_color = (int(color[i][0]), int(color[i][1]), int(color[i][2]))
cv2.line(
display,
(left_x, left_y),
(right_x, right_y),
cur_color,
1,
lineType=cv2.LINE_AA,
)
return display
|