|
|
import argparse |
|
|
import os |
|
|
import sys |
|
|
import time |
|
|
import warnings |
|
|
from importlib import import_module |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from PIL import Image |
|
|
|
|
|
warnings.filterwarnings("ignore", category=UserWarning, module="torch.nn.functional") |
|
|
|
|
|
|
|
|
def str2bool(v: str, strict=True) -> bool: |
|
|
if isinstance(v, bool): |
|
|
return v |
|
|
elif isinstance(v, str): |
|
|
if v.lower() in ("true", "yes", "on" "t", "y", "1"): |
|
|
return True |
|
|
elif v.lower() in ("false", "no", "off", "f", "n", "0"): |
|
|
return False |
|
|
if strict: |
|
|
raise argparse.ArgumentTypeError("Unsupported value encountered.") |
|
|
else: |
|
|
return True |
|
|
|
|
|
|
|
|
def to_cuda(data, device="cuda", exclude_keys: "list[str]" = None): |
|
|
if isinstance(data, torch.Tensor): |
|
|
data = data.to(device) |
|
|
elif isinstance(data, (tuple, list, set)): |
|
|
data = [to_cuda(b, device) for b in data] |
|
|
elif isinstance(data, dict): |
|
|
if exclude_keys is None: |
|
|
exclude_keys = [] |
|
|
for k in data.keys(): |
|
|
if k not in exclude_keys: |
|
|
data[k] = to_cuda(data[k], device) |
|
|
else: |
|
|
|
|
|
data = data |
|
|
return data |
|
|
|
|
|
|
|
|
class HiddenPrints: |
|
|
def __enter__(self): |
|
|
self._original_stdout = sys.stdout |
|
|
sys.stdout = open(os.devnull, "w") |
|
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb): |
|
|
sys.stdout.close() |
|
|
sys.stdout = self._original_stdout |
|
|
|
|
|
|
|
|
class Logger(object): |
|
|
def __init__(self): |
|
|
self.terminal = sys.stdout |
|
|
self.file = None |
|
|
|
|
|
def open(self, file, mode=None): |
|
|
if mode is None: |
|
|
mode = "w" |
|
|
self.file = open(file, mode) |
|
|
|
|
|
def write(self, message, is_terminal=1, is_file=1): |
|
|
if "\r" in message: |
|
|
is_file = 0 |
|
|
if is_terminal == 1: |
|
|
self.terminal.write(message) |
|
|
self.terminal.flush() |
|
|
if is_file == 1: |
|
|
self.file.write(message) |
|
|
self.file.flush() |
|
|
|
|
|
def flush(self): |
|
|
|
|
|
|
|
|
|
|
|
pass |
|
|
|
|
|
|
|
|
def get_network(arch: str, isTrain=False, continue_train=False, init_gain=0.02, pretrained=True): |
|
|
if "resnet" in arch: |
|
|
from networks.resnet import ResNet |
|
|
|
|
|
resnet = getattr(import_module("networks.resnet"), arch) |
|
|
if isTrain: |
|
|
if continue_train: |
|
|
model: ResNet = resnet(num_classes=1) |
|
|
else: |
|
|
model: ResNet = resnet(pretrained=pretrained) |
|
|
model.fc = nn.Linear(2048, 1) |
|
|
nn.init.normal_(model.fc.weight.data, 0.0, init_gain) |
|
|
else: |
|
|
model: ResNet = resnet(num_classes=1) |
|
|
return model |
|
|
else: |
|
|
raise ValueError(f"Unsupported arch: {arch}") |
|
|
|
|
|
|
|
|
def pad_img_to_square(img: np.ndarray): |
|
|
H, W = img.shape[:2] |
|
|
if H != W: |
|
|
new_size = max(H, W) |
|
|
img = np.pad(img, ((0, new_size - H), (0, new_size - W), (0, 0)), mode="constant") |
|
|
assert img.shape[0] == img.shape[1] == new_size |
|
|
return img |
|
|
|
