File size: 2,062 Bytes
62456b0 |
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 |
"""
Different utilities such as orthogonalization of weights, initialization of
loggers, etc
Copyright (C) 2018, Matias Tassano <matias.tassano@parisdescartes.fr>
This program is free software: you can use, modify and/or
redistribute it under the terms of the GNU General Public
License as published by the Free Software Foundation, either
version 3 of the License, or (at your option) any later
version. You should have received a copy of this license along
this program. If not, see <http://www.gnu.org/licenses/>.
"""
import numpy as np
import cv2
def variable_to_cv2_image(varim):
r"""Converts a torch.autograd.Variable to an OpenCV image
Args:
varim: a torch.autograd.Variable
"""
nchannels = varim.size()[1]
if nchannels == 1:
res = (varim.data.cpu().numpy()[0, 0, :]*255.).clip(0, 255).astype(np.uint8)
elif nchannels == 3:
res = varim.data.cpu().numpy()[0]
res = cv2.cvtColor(res.transpose(1, 2, 0), cv2.COLOR_RGB2BGR)
res = (res*255.).clip(0, 255).astype(np.uint8)
else:
raise Exception('Number of color channels not supported')
return res
def normalize(data):
return np.float32(data/255.)
def remove_dataparallel_wrapper(state_dict):
r"""Converts a DataParallel model to a normal one by removing the "module."
wrapper in the module dictionary
Args:
state_dict: a torch.nn.DataParallel state dictionary
"""
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, vl in state_dict.items():
name = k[7:] # remove 'module.' of DataParallel
new_state_dict[name] = vl
return new_state_dict
def is_rgb(im_path):
r""" Returns True if the image in im_path is an RGB image
"""
from skimage.io import imread
rgb = False
im = imread(im_path)
if (len(im.shape) == 3):
if not(np.allclose(im[...,0], im[...,1]) and np.allclose(im[...,2], im[...,1])):
rgb = True
print("rgb: {}".format(rgb))
print("im shape: {}".format(im.shape))
return rgb
|