|
|
import logging |
|
|
import os |
|
|
import torch |
|
|
import shutil |
|
|
from torchvision import transforms |
|
|
import numpy as np |
|
|
import random |
|
|
import cv2 |
|
|
|
|
|
|
|
|
class Logger(): |
|
|
def __init__(self, path="log.txt"): |
|
|
self.logger = logging.getLogger('DCFM') |
|
|
self.file_handler = logging.FileHandler(path, "w") |
|
|
self.stdout_handler = logging.StreamHandler() |
|
|
self.stdout_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s')) |
|
|
self.file_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s')) |
|
|
self.logger.addHandler(self.file_handler) |
|
|
self.logger.addHandler(self.stdout_handler) |
|
|
self.logger.setLevel(logging.INFO) |
|
|
self.logger.propagate = False |
|
|
|
|
|
def info(self, txt): |
|
|
self.logger.info(txt) |
|
|
|
|
|
def close(self): |
|
|
self.file_handler.close() |
|
|
self.stdout_handler.close() |
|
|
|
|
|
class AverageMeter(object): |
|
|
"""Computes and stores the average and current value""" |
|
|
def __init__(self): |
|
|
self.reset() |
|
|
|
|
|
def reset(self): |
|
|
self.val = 0.0 |
|
|
self.avg = 0.0 |
|
|
self.sum = 0.0 |
|
|
self.count = 0.0 |
|
|
|
|
|
def update(self, val, n=1): |
|
|
self.val = val |
|
|
self.sum += val * n |
|
|
self.count += n |
|
|
self.avg = self.sum / self.count |
|
|
|
|
|
|
|
|
def save_checkpoint(state, path, filename="checkpoint.pth"): |
|
|
torch.save(state, os.path.join(path, filename)) |
|
|
|
|
|
|
|
|
def save_tensor_img(tenor_im, path): |
|
|
im = tenor_im.cpu().clone() |
|
|
im = im.squeeze(0) |
|
|
tensor2pil = transforms.ToPILImage() |
|
|
im = tensor2pil(im) |
|
|
im.save(path) |
|
|
|
|
|
|
|
|
def save_tensor_merge(tenor_im, tensor_mask, path, colormap='HOT'): |
|
|
im = tenor_im.cpu().detach().clone() |
|
|
im = im.squeeze(0).numpy() |
|
|
im = ((im - np.min(im)) / (np.max(im) - np.min(im) + 1e-20)) * 255 |
|
|
im = np.array(im,np.uint8) |
|
|
mask = tensor_mask.cpu().detach().clone() |
|
|
mask = mask.squeeze(0).numpy() |
|
|
mask = ((mask - np.min(mask)) / (np.max(mask) - np.min(mask) + 1e-20)) * 255 |
|
|
mask = np.clip(mask, 0, 255) |
|
|
mask = np.array(mask, np.uint8) |
|
|
if colormap == 'HOT': |
|
|
mask = cv2.applyColorMap(mask[0,:,:], cv2.COLORMAP_HOT) |
|
|
elif colormap == 'PINK': |
|
|
mask = cv2.applyColorMap(mask[0,:,:], cv2.COLORMAP_PINK) |
|
|
elif colormap == 'BONE': |
|
|
mask = cv2.applyColorMap(mask[0,:,:], cv2.COLORMAP_BONE) |
|
|
|
|
|
im = im.transpose((1, 2, 0)) |
|
|
im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR) |
|
|
mix = cv2.addWeighted(im, 0.3, mask, 0.7, 0) |
|
|
cv2.imwrite(path, mix) |
|
|
|
|
|
|
|
|
def set_seed(seed): |
|
|
torch.manual_seed(seed) |
|
|
torch.cuda.manual_seed(seed) |
|
|
torch.cuda.manual_seed_all(seed) |
|
|
np.random.seed(seed) |
|
|
random.seed(seed) |
|
|
torch.backends.cudnn.deterministic = True |
|
|
torch.backends.cudnn.benchmark = False |
|
|
|
|
|
|
|
|
|
