File size: 2,883 Bytes
f7f604d |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import re
import os
import yaml
import cv2
import argparse
import warnings
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from easydict import EasyDict as ed
class Simplify(nn.Module):
def __init__(self, model):
super(Simplify, self).__init__()
self.model = model
def cuda(self):
self.model = self.model.cuda()
return self
def forward(self, x):
out = self.model({'image': x})
return out['pred']
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--config', '-c', type=str, default='configs/InSPyReNet_SwinB.yaml')
parser.add_argument('--resume', '-r', action='store_true', default=False)
parser.add_argument('--verbose', '-v', action='store_true', default=False)
parser.add_argument('--debug', '-d', action='store_true', default=False)
args = parser.parse_args()
cuda_visible_devices = None
local_rank = -1
if "CUDA_VISIBLE_DEVICES" in os.environ.keys():
cuda_visible_devices = [int(i) for i in os.environ["CUDA_VISIBLE_DEVICES"].split(',')]
if "LOCAL_RANK" in os.environ.keys():
local_rank = int(os.environ["LOCAL_RANK"])
if local_rank == -1:
device_num = 1
elif cuda_visible_devices is None:
device_num = torch.cuda.device_count()
else:
device_num = len(cuda_visible_devices)
args.device_num = device_num
args.local_rank = local_rank
warnings.simplefilter("ignore")
return args
def sort(x):
convert = lambda text: int(text) if text.isdigit() else text.lower()
alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
return sorted(x, key=alphanum_key)
def load_config(config_dir, easy=True):
cfg = yaml.load(open(config_dir), yaml.FullLoader)
if easy is True:
cfg = ed(cfg)
return cfg
def to_cuda(sample):
for key in sample.keys():
if type(sample[key]) == torch.Tensor:
sample[key] = sample[key].cuda()
return sample
def to_numpy(pred, shape):
pred = F.interpolate(pred, shape, mode='bilinear', align_corners=True)
pred = pred.data.cpu()
pred = pred.numpy().squeeze()
return pred
def debug_tile(deblist, size=(100, 100), activation=None):
debugs = []
for debs in deblist:
debug = []
for deb in debs:
if activation is not None:
deb = activation(deb)
log = deb.cpu().detach().numpy().squeeze()
log = ((log - log.min()) / (log.max() - log.min()) * 255).astype(np.uint8)
log = cv2.cvtColor(log, cv2.COLOR_GRAY2RGB)
log = cv2.resize(log, size)
debug.append(log)
debugs.append(np.vstack(debug))
return np.hstack(debugs)
if __name__ == "__main__":
x = torch.rand(4, 3, 576, 576) |