Spaces:
Runtime error
Runtime error
File size: 1,717 Bytes
1f418ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import torch
import cv2
import numpy as np
import torch.nn.functional as F
def get_activation(name, bank):
def hook(model, input, output):
bank[name] = output
return hook
class HookTool:
def __init__(self):
self.feat = None
def hook_in_fun(self, module, fea_in, fea_out):
self.feat = fea_in
def hook_out_fun(self, module, fea_in, fea_out):
self.feat = fea_out
class RunningAverageMap:
""" Saving avg depth estimation results."""
def __init__(self, average_map, count_map):
self.average_map = average_map
self.count_map = count_map
self.average_map = self.average_map / self.count_map
def update(self, pred_map, ct_map):
self.average_map = (pred_map + self.count_map * self.average_map) / (self.count_map + ct_map)
self.count_map = self.count_map + ct_map
def resize(self, resolution):
temp_avg_map = self.average_map.unsqueeze(0).unsqueeze(0)
temp_count_map = self.count_map.unsqueeze(0).unsqueeze(0)
self.average_map = F.interpolate(temp_avg_map, size=resolution).squeeze()
self.count_map = F.interpolate(temp_count_map, size=resolution, mode='bilinear', align_corners=True).squeeze()
def generatemask(size):
# Generates a Guassian mask
mask = np.zeros(size, dtype=np.float32)
sigma = int(size[0]/16)
k_size = int(2 * np.ceil(2 * int(size[0]/16)) + 1)
mask[int(0.1*size[0]):size[0] - int(0.1*size[0]), int(0.1*size[1]): size[1] - int(0.1*size[1])] = 1
mask = cv2.GaussianBlur(mask, (int(k_size), int(k_size)), sigma)
mask = (mask - mask.min()) / (mask.max() - mask.min())
mask = mask.astype(np.float32)
return mask |