rename py/BiRefNet to py/BiRefNet_v2, for avoid module name conflicts

py/BiRefNet_v2/utils.py

@@ -0,0 +1,97 @@
+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