98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
import logging
|
|
import os
|
|
import torch
|
|
from torchvision import transforms
|
|
import numpy as np
|
|
import random
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
|
|
def path_to_image(path, size=(1024, 1024), color_type=['rgb', 'gray'][0]):
|
|
if color_type.lower() == 'rgb':
|
|
image = cv2.imread(path)
|
|
elif color_type.lower() == 'gray':
|
|
image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
|
|
else:
|
|
print('Select the color_type to return, either to RGB or gray image.')
|
|
return
|
|
if size:
|
|
image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
|
|
if color_type.lower() == 'rgb':
|
|
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert('RGB')
|
|
else:
|
|
image = Image.fromarray(image).convert('L')
|
|
return image
|
|
|
|
|
|
|
|
def check_state_dict(state_dict, unwanted_prefix='_orig_mod.'):
|
|
for k, v in list(state_dict.items()):
|
|
if k.startswith(unwanted_prefix):
|
|
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
|
|
return state_dict
|
|
|
|
|
|
def generate_smoothed_gt(gts):
|
|
epsilon = 0.001
|
|
new_gts = (1-epsilon)*gts+epsilon/2
|
|
return new_gts
|
|
|
|
|
|
class Logger():
|
|
def __init__(self, path="log.txt"):
|
|
self.logger = logging.getLogger('BiRefNet')
|
|
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="latest.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 set_seed(seed):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
torch.backends.cudnn.deterministic = True
|