WSCL / utils /misc.py
yhzhai's picture
release code
482ab8a
import copy
import datetime
import json
import math
import os
import random
import signal
import subprocess
import sys
import time
import warnings
from collections import defaultdict
from shutil import copy2
from typing import Dict
import numpy as np
import prettytable as pt
import torch
import torch.nn as nn
from termcolor import cprint
from torch.utils.tensorboard import SummaryWriter
class Logger(object):
def __init__(self, filename, stream=sys.stdout):
self.terminal = stream
self.log = open(filename, "a")
def write(self, message):
self.terminal.write(message)
self.log.write(message)
def flush(self):
pass
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.sum = 0
self.avg = 0
self.val = 0
self.count = 0
def reset(self):
self.sum = 0
self.avg = 0
self.val = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum = self.sum + val * n
self.count = self.count + n
self.avg = self.sum / self.count
def __str__(self):
return f"{self.avg: .5f}"
def get_sha():
"""Get git current status"""
cwd = os.path.dirname(os.path.abspath(__file__))
def _run(command):
return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
sha = "N/A"
diff = "clean"
branch = "N/A"
message = "N/A"
try:
sha = _run(["git", "rev-parse", "HEAD"])
sha = sha[:8]
subprocess.check_output(["git", "diff"], cwd=cwd)
diff = _run(["git", "diff-index", "HEAD"])
diff = "has uncommited changes" if diff else "clean"
branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
message = _run(["git", "log", "--pretty=format:'%s'", sha, "-1"]).replace(
"'", ""
)
except Exception:
pass
return {"sha": sha, "status": diff, "branch": branch, "prev_commit": message}
def setup_env(opt):
if opt.eval or opt.debug:
opt.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.autograd.set_detect_anomaly(True)
return None
dir_name = opt.dir_name
save_root_path = opt.save_root_path
if not os.path.exists(save_root_path):
os.mkdir(save_root_path)
# deterministic
torch.manual_seed(opt.seed)
np.random.seed(opt.seed)
random.seed(opt.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
# mkdir subdirectories
checkpoint = "checkpoint"
if not os.path.exists(os.path.join(save_root_path, dir_name)):
os.mkdir(os.path.join(save_root_path, dir_name))
os.mkdir(os.path.join(save_root_path, dir_name, checkpoint))
# save log
sys.stdout = Logger(os.path.join(save_root_path, dir_name, "log.log"), sys.stdout)
sys.stderr = Logger(os.path.join(save_root_path, dir_name, "error.log"), sys.stderr)
# save parameters
params = copy.deepcopy(vars(opt))
params.pop("device")
with open(os.path.join(save_root_path, dir_name, "params.json"), "w") as f:
json.dump(params, f)
# print info
print(
"Running on {}, PyTorch version {}, files will be saved at {}".format(
opt.device, torch.__version__, os.path.join(save_root_path, dir_name)
)
)
print("Devices:")
for i in range(torch.cuda.device_count()):
print(" {}:".format(i), torch.cuda.get_device_name(i))
print(f"Git: {get_sha()}.")
# return tensorboard summarywriter
return SummaryWriter("{}/{}/".format(opt.save_root_path, opt.dir_name))
class MetricLogger(object):
def __init__(self, delimiter=" ", writer=None, suffix=None):
self.meters = defaultdict(AverageMeter)
self.delimiter = delimiter
self.writer = writer
self.suffix = suffix
def update(self, **kwargs):
for k, v in kwargs.items():
if isinstance(v, torch.Tensor):
v = v.item()
assert isinstance(v, (float, int)), f"Unsupport type {type(v)}."
self.meters[k].update(v)
def add_meter(self, name, meter):
self.meters[name] = meter
def get_meters(self):
result = {}
for k, v in self.meters.items():
result[k] = v.avg
return result
def prepend_subprefix(self, subprefix: str):
old_keys = list(self.meters.keys())
for k in old_keys:
self.meters[k.replace("/", f"/{subprefix}")] = self.meters[k]
for k in old_keys:
del self.meters[k]
def log_every(self, iterable, print_freq=10, header=""):
i = 0
start_time = time.time()
end = time.time()
iter_time = AverageMeter()
space_fmt = ":" + str(len(str(len(iterable)))) + "d"
log_msg = self.delimiter.join(
[
header,
"[{0" + space_fmt + "}/{1}]",
"eta: {eta}",
"{meters}",
"iter time: {time}s",
]
)
for obj in iterable:
yield i, obj
iter_time.update(time.time() - end)
if (i + 1) % print_freq == 0 or i == len(iterable) - 1:
eta_seconds = iter_time.avg * (len(iterable) - i)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
print(
log_msg.format(
i + 1,
len(iterable),
eta=eta_string,
meters=str(self),
time=str(iter_time),
).replace(" ", " ")
)
i += 1
end = time.time()
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print(
"{} Total time: {} ({:.4f}s / it)".format(
header, total_time_str, total_time / len(iterable)
)
)
def write_tensorboard(self, step):
if self.writer is not None:
for k, v in self.meters.items():
# if self.suffix:
# self.writer.add_scalar(
# '{}/{}'.format(k, self.suffix), v.avg, step)
# else:
self.writer.add_scalar(k, v.avg, step)
def stat_table(self):
tb = pt.PrettyTable(field_names=["Metrics", "Values"])
for name, meter in self.meters.items():
tb.add_row([name, str(meter)])
return tb.get_string()
def __getattr__(self, attr):
if attr in self.meters:
return self.meters[attr]
if attr in self.__dict__:
return self.__dict__[attr]
raise AttributeError(
"'{}' object has no attribute '{}'".format(type(self).__name__, attr)
)
def __str__(self):
loss_str = []
for name, meter in self.meters.items():
loss_str.append("{}: {}".format(name, str(meter)))
return self.delimiter.join(loss_str).replace(" ", " ")
def save_model(path, model: nn.Module, epoch, opt, performance=None):
if not opt.debug:
try:
torch.save(
{
"model": model.state_dict(),
"epoch": epoch,
"opt": opt,
"performance": performance,
},
path,
)
except Exception as e:
cprint("Failed to save {} because {}".format(path, str(e)))
def resume_from(model: nn.Module, resume_path: str):
checkpoint = torch.load(resume_path, map_location="cpu")
state_dict = checkpoint["model"]
performance = checkpoint["performance"]
try:
model.load_state_dict(state_dict)
except Exception as e:
model.load_state_dict(state_dict, strict=False)
cprint("Failed to load full model because {}".format(str(e)), "red")
time.sleep(3)
print(f"{resume_path} model loaded. It performance is")
if performance is not None:
for k, v in performance.items():
print(f"{k}: {v}")
def update_record(result: Dict, epoch: int, opt, file_name: str = "latest_record"):
if not opt.debug:
# save txt file
tb = pt.PrettyTable(field_names=["Metrics", "Values"])
with open(
os.path.join(opt.save_root_path, opt.dir_name, f"{file_name}.txt"), "w"
) as f:
f.write(f"Performance at {epoch}-th epoch:\n\n")
for k, v in result.items():
tb.add_row([k, "{:.7f}".format(v)])
f.write(tb.get_string())
# save json file
result["epoch"] = epoch
with open(
os.path.join(opt.save_root_path, opt.dir_name, f"{file_name}.json"), "w"
) as f:
json.dump(result, f)
def pixel_acc(pred, label):
"""Compute pixel-level prediction accuracy."""
warnings.warn("I am not sure if this implementation is correct.")
label_size = label.shape[-2:]
if pred.shape[-2] != label_size:
pred = torch.nn.functional.interpolate(
pred, size=label_size, mode="bilinear", align_corners=False
)
pred[torch.where(pred > 0.5)] = 1
pred[torch.where(pred <= 0.5)] = 0
correct = torch.sum((pred + label) == 1.0)
total = torch.numel(pred)
return correct / (total + 1e-8)
def calculate_pixel_f1(pd, gt, prefix="", suffix=""):
if np.max(pd) == np.max(gt) and np.max(pd) == 0:
f1, iou = 1.0, 1.0
return f1, 0.0, 0.0
seg_inv, gt_inv = np.logical_not(pd), np.logical_not(gt)
true_pos = float(np.logical_and(pd, gt).sum())
false_pos = np.logical_and(pd, gt_inv).sum()
false_neg = np.logical_and(seg_inv, gt).sum()
f1 = 2 * true_pos / (2 * true_pos + false_pos + false_neg + 1e-6)
precision = true_pos / (true_pos + false_pos + 1e-6)
recall = true_pos / (true_pos + false_neg + 1e-6)
return {
f"{prefix}pixel_f1{suffix}": f1,
f"{prefix}pixel_prec{suffix}": precision,
f"{prefix}pixel_recall{suffix}": recall,
}
def calculate_img_score(pd, gt, prefix="", suffix="", eta=1e-6):
seg_inv, gt_inv = np.logical_not(pd), np.logical_not(gt)
true_pos = float(np.logical_and(pd, gt).sum())
false_pos = float(np.logical_and(pd, gt_inv).sum())
false_neg = float(np.logical_and(seg_inv, gt).sum())
true_neg = float(np.logical_and(seg_inv, gt_inv).sum())
acc = (true_pos + true_neg) / (true_pos + true_neg + false_neg + false_pos + eta)
sen = true_pos / (true_pos + false_neg + eta)
spe = true_neg / (true_neg + false_pos + eta)
precision = true_pos / (true_pos + false_pos + eta)
recall = true_pos / (true_pos + false_neg + eta)
try:
f1 = 2 * sen * spe / (sen + spe)
except:
f1 = -math.inf
return {
f"{prefix}image_acc{suffix}": acc,
f"{prefix}image_sen{suffix}": sen,
f"{prefix}image_spe{suffix}": spe,
f"{prefix}image_f1{suffix}": f1,
f"{prefix}image_true_pos{suffix}": true_pos,
f"{prefix}image_true_neg{suffix}": true_neg,
f"{prefix}image_false_pos{suffix}": false_pos,
f"{prefix}image_false_neg{suffix}": false_neg,
f"{prefix}image_prec{suffix}": precision,
f"{prefix}image_recall{suffix}": recall,
}
class timeout:
def __init__(self, seconds=1, error_message="Timeout"):
self.seconds = seconds
self.error_message = error_message
def handle_timeout(self, signum, frame):
raise TimeoutError(self.error_message)
def __enter__(self):
signal.signal(signal.SIGALRM, self.handle_timeout)
signal.alarm(self.seconds)
def __exit__(self, type, value, traceback):
signal.alarm(0)