core/utils/utils.py

import os
import sys
import time
import shutil
import logging
import numpy as np

from scipy import interpolate
from datetime import datetime

import torch
import torch.nn.functional as F



class InputPadder:
    """ Pads images such that dimensions are divisible by 8 """
    def __init__(self, dims, mode='sintel', divis_by=8):
        self.ht, self.wd = dims[-2:]
        pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by
        pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by
        if mode == 'sintel':
            self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
        else:
            self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]

    def pad(self, *inputs):
        assert all((x.ndim == 4) for x in inputs)
        return [F.pad(x, self._pad, mode='replicate') for x in inputs]
    
    def pad_intrinsics(self, intrinsic):
        intrinsic[:, 2] += self._pad[0]
        intrinsic[:, 3] += self._pad[2]
        return intrinsic

    def unpad(self, x):
        assert x.ndim == 4
        ht, wd = x.shape[-2:]
        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
        return x[..., c[0]:c[1], c[2]:c[3]]

def forward_interpolate(flow):
    flow = flow.detach().cpu().numpy()
    dx, dy = flow[0], flow[1]

    ht, wd = dx.shape
    x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))

    x1 = x0 + dx
    y1 = y0 + dy
    
    x1 = x1.reshape(-1)
    y1 = y1.reshape(-1)
    dx = dx.reshape(-1)
    dy = dy.reshape(-1)

    valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
    x1 = x1[valid]
    y1 = y1[valid]
    dx = dx[valid]
    dy = dy[valid]

    flow_x = interpolate.griddata(
        (x1, y1), dx, (x0, y0), method='nearest', fill_value=0)

    flow_y = interpolate.griddata(
        (x1, y1), dy, (x0, y0), method='nearest', fill_value=0)

    flow = np.stack([flow_x, flow_y], axis=0)
    return torch.from_numpy(flow).float()


def bilinear_sampler(img, coords, mode='bilinear', mask=False):
    """ Wrapper for grid_sample, uses pixel coordinates """
    H, W = img.shape[-2:]
    xgrid, ygrid = coords.split([1,1], dim=-1)
    xgrid = 2*xgrid/(W-1) - 1
    if H > 1:
        ygrid = 2*ygrid/(H-1) - 1

    grid = torch.cat([xgrid, ygrid], dim=-1)
    img = F.grid_sample(img, grid, align_corners=True)

    if mask:
        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
        return img, mask.float()

    return img


def coords_grid(batch, ht, wd):
    coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
    coords = torch.stack(coords[::-1], dim=0).float()
    return coords[None].repeat(batch, 1, 1, 1)

def hor_coords_grid(batch, ht, wd):
    # (batch,1,H,W)
    hor_coords = torch.arange(wd).float().repeat(batch, 1, ht, 1)
    return hor_coords


def upflow8(flow, mode='bilinear'):
    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
    return  8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)

def gauss_blur(input, N=5, std=1):
    B, D, H, W = input.shape
    x, y = torch.meshgrid(torch.arange(N).float() - N//2, torch.arange(N).float() - N//2)
    unnormalized_gaussian = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * std ** 2))
    weights = unnormalized_gaussian / unnormalized_gaussian.sum().clamp(min=1e-4)
    weights = weights.view(1,1,N,N).to(input)
    output = F.conv2d(input.reshape(B*D,1,H,W), weights, padding=N//2)
    return output.view(B, D, H, W)

def disparity_computation(params, slant=None, slant_norm=False, coords0=None):
    """
    args:
        params: (B,C,...), C is the type of parameters.
        coords0: (B,C,...), C is the number of coordinates' axis.
    """
    if slant is None or len(slant)==0 :
        offset = params
    elif slant=="slant" :
        # d = a*u + b*v + c
        B,H,W = coords0.shape[0], coords0.shape[-2], coords0.shape[-1]
        if slant_norm:
            norm_range = torch.Tensor([W,H])[None,:,None,None].float().to(coords0.device)
            offset = params[:,0] * coords0[:,0] / norm_range[:,0] + \
                     params[:,1] * coords0[:,1] / norm_range[:,1] + \
                     params[:,2]
        else:
            offset = params[:,0] * coords0[:,0] + \
                     params[:,1] * coords0[:,1] + \
                     params[:,2]
    elif slant=="slant_local":
        raise Exception("slant_local is not supported")
    else:
        raise Exception(f"{slant} is not supported")
    return offset


def sv_intermediate_results(data, name, sv_path):
    try:
        sv_path = os.path.join(sv_path, "data")
        if not os.path.exists(sv_path):
            os.makedirs(sv_path)
        
        data_numpy = data.cpu().data.numpy()
        np.save(os.path.join(sv_path, name+".npy"), data_numpy)
        # print("saving to {}".format( os.path.join(sv_path, name+".npy") ))
    except Exception as err:
        raise Exception(err, data.shape, name, sv_path)

def load_intermediate_results(name, sv_path):
    sv_path = os.path.join(sv_path, "data")
    data = np.load(os.path.join(sv_path, name+".npy"))
    return data


def rescale_modulation(itr, iters, modulation_alg, modulation_ratio):
    # we hope modulation has less effect at the first several iterations as the disp is unreliable and the lcoal LBP disp is unreliable
    if modulation_alg == "linear":
        ratio = modulation_ratio * itr / iters
    elif modulation_alg == "sigmoid":
        ratio = modulation_ratio * 1 / (1 + np.exp(-2 * (itr - 5)))
    else:
        raise Exception("Not supported modulation_alg: {}".format(modulation_alg))
    return ratio



NODE_RANK    = os.getenv('NODE_RANK', default=0)
LOCAL_RANK   = os.getenv("LOCAL_RANK", default=0)
LOG_ROOT     = os.getenv('LOG_ROOT', default="logs")
TB_ROOT      = os.getenv('TB_ROOT', default="runs")

class LoggerCommon:
    def __init__(self, name):
        self.name = name
        self.log_name = '{}-{}.log'.format(name, datetime.now().strftime("%y%m%d_%H%M%S"))
        self.log_path = os.path.join(LOG_ROOT, self.log_name)
        if int(LOCAL_RANK)==0 and int(NODE_RANK)==0:
            os.makedirs(LOG_ROOT, exist_ok=True)
            logging.basicConfig(level=logging.INFO,
                                format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s',
                                handlers = [logging.FileHandler(self.log_path), 
                                            logging.StreamHandler()]
                            )
            self.logger = logging.getLogger(name)
            self.logger.addHandler(logging.FileHandler(self.log_path))
    
    def _set_handlers(self):
        # 清除之前的所有处理器
        self.logger.handlers.clear()
        
        # 设置文件和控制台处理器
        file_handler = logging.FileHandler(self.log_path)
        console_handler = logging.StreamHandler()

        # 设置处理器格式
        formatter = logging.Formatter('%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s')
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)

        # 添加处理器
        self.logger.addHandler(file_handler)
        self.logger.addHandler(console_handler)

    def set_log_path(self, new_log_path, name=None):
        # 删除旧日志文件（如果存在）
        if os.path.exists(self.log_path):
            os.remove(self.log_path)
        
        # 更新路径并重设处理器
        if name is not None:
            self.name = name
        self.log_name = '{}-{}.log'.format(self.name, datetime.now().strftime("%y%m%d_%H%M%S"))
        self.log_path = os.path.join(new_log_path, self.log_name)
        os.makedirs(new_log_path, exist_ok=True)
        self._set_handlers()

    
    def info(self, message):
        if int(LOCAL_RANK)==0 and int(NODE_RANK)==0:
            self.logger.info(message)
    
    def warning(self, message):
        if int(LOCAL_RANK)==0 and int(NODE_RANK)==0:
            self.logger.warning(message)
    
    def error(self, message):
        if int(LOCAL_RANK)==0 and int(NODE_RANK)==0:
            self.logger.error(message)
    
    def exception(self, message):
        if int(LOCAL_RANK)==0 and int(NODE_RANK)==0:
            self.logger.exception(message)
    
    def print_args(self, args):
        msg = ""
        args_dict = vars(args)
        max_arg_length = max(len(arg_name) for arg_name in args_dict.keys())
        for arg_name, arg_value in args_dict.items():
            arg_name_padded = arg_name.ljust(max_arg_length)
            msg += f"{arg_name_padded}: {arg_value}\r\n"
        self.info(msg)
        


from torch.utils.tensorboard import SummaryWriter

class LoggerTraining(LoggerCommon):

    SUM_FREQ = 100

    def __init__(self, name, model=None, scheduler=None):
        super(LoggerTraining, self).__init__(name)

        if int(LOCAL_RANK)==0 and int(NODE_RANK)==0:
            os.makedirs(TB_ROOT, exist_ok=True)

        self.model = model
        self.scheduler = scheduler
        self.silence = False
        self.total_steps = 0
        self.running_loss = {}
        self.writer = SummaryWriter(log_dir=TB_ROOT)
    
    def set_training(self, model, scheduler):
        self.model = model
        self.scheduler = scheduler
    
    def _print_training_status(self):
        metrics_data = [self.running_loss[k]/LoggerTraining.SUM_FREQ for k in sorted(self.running_loss.keys())]
        training_str = "[{:6d}, {:10.7f}] ".format(self.total_steps+1, self.scheduler.get_last_lr()[0])
        metrics_str = ("{:10.4f}, "*len(metrics_data)).format(*metrics_data)
        
        # print the training status
        self.info(f"Training Metrics ({self.total_steps}): {training_str + metrics_str}")

        if self.writer is None:
            self.writer = SummaryWriter(log_dir=TB_ROOT)

        for k in self.running_loss:
            self.writer.add_scalar(k, self.running_loss[k]/LoggerTraining.SUM_FREQ, self.total_steps)
            self.running_loss[k] = 0.0

    def push(self, metrics):
        self.total_steps += 1

        for key in metrics:
            if key not in self.running_loss:
                self.running_loss[key] = 0.0

            self.running_loss[key] += metrics[key]

        if self.total_steps % LoggerTraining.SUM_FREQ == LoggerTraining.SUM_FREQ-1:
            self._print_training_status()
            self.running_loss = {}

    def write_dict(self, results):
        if self.writer is None:
            self.writer = SummaryWriter(log_dir=TB_ROOT)

        for key in results:
            self.writer.add_scalar(key, results[key], self.total_steps)

    def close(self):
        self.writer.close()



def init_directories(directories):
    if int(LOCAL_RANK)==0 and int(NODE_RANK)==0 :
        for directory in directories:
            os.makedirs(directory, exist_ok=True)
    
def delete_directories_if_static(directories):
    if int(LOCAL_RANK)==0 and int(NODE_RANK)==0 :
        # 如果检测到文件大小有变化，终止删除操作
        if not is_any_folder_static(directories):
            print("File sizes are changing in one of the directories {}.".format(directories) + \
                  "No directories will be deleted.")
            return
        
        # 如果所有文件都静止，删除目录
        for directory in directories:
            if os.path.exists(directory):
                shutil.rmtree(directory)
                print(f"Directory {directory} deleted")

def get_file_sizes(directories):
    """返回多个目录中所有文件的大小字典"""
    file_sizes = {}
    for directory in directories:
        if os.path.exists(directory):
            for root, dirs, files in os.walk(directory):
                for file in files:
                    filepath = os.path.join(root, file)
                    file_sizes[filepath] = os.path.getsize(filepath)
    return file_sizes

def is_any_folder_static(directories, check_interval=2):
    """检测所有文件是否静止（没有变化）"""
    # 获取所有文件初始大小
    initial_sizes = get_file_sizes(directories)
    time.sleep(check_interval)  # 等待一段时间，观察文件变化
    final_sizes = {filepath: os.path.getsize(filepath) for filepath in initial_sizes if os.path.exists(filepath)}
    
    # 如果文件大小一致，则所有文件静止
    return initial_sizes == final_sizes