Spaces:
Sleeping
Sleeping
File size: 3,136 Bytes
34d43ef |
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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
#!/usr/bin/python
#
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import PIL
import torch
import torchvision.transforms as T
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
INV_IMAGENET_MEAN = [-m for m in IMAGENET_MEAN]
INV_IMAGENET_STD = [1.0 / s for s in IMAGENET_STD]
def imagenet_preprocess():
return T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
def rescale(x):
lo, hi = x.min(), x.max()
return x.sub(lo).div(hi - lo)
def imagenet_deprocess(rescale_image=True):
transforms = [
T.Normalize(mean=[0, 0, 0], std=INV_IMAGENET_STD),
T.Normalize(mean=INV_IMAGENET_MEAN, std=[1.0, 1.0, 1.0]),
]
if rescale_image:
transforms.append(rescale)
return T.Compose(transforms)
def imagenet_deprocess_batch(imgs, rescale=True):
"""
Input:
- imgs: FloatTensor of shape (N, C, H, W) giving preprocessed images
Output:
- imgs_de: ByteTensor of shape (N, C, H, W) giving deprocessed images
in the range [0, 255]
"""
if isinstance(imgs, torch.autograd.Variable):
imgs = imgs.data
imgs = imgs.cpu().clone()
deprocess_fn = imagenet_deprocess(rescale_image=rescale)
imgs_de = []
for i in range(imgs.size(0)):
img_de = deprocess_fn(imgs[i])[None]
img_de = img_de.mul(255).clamp(0, 255).byte()
imgs_de.append(img_de)
imgs_de = torch.cat(imgs_de, dim=0)
return imgs_de
class Resize(object):
def __init__(self, size, interp=PIL.Image.BILINEAR):
if isinstance(size, tuple):
H, W = size
self.size = (W, H)
else:
self.size = (size, size)
self.interp = interp
def __call__(self, img):
return img.resize(self.size, self.interp)
def unpack_var(v):
if isinstance(v, torch.autograd.Variable):
return v.data
return v
def split_graph_batch(triples, obj_data, obj_to_img, triple_to_img):
triples = unpack_var(triples)
obj_data = [unpack_var(o) for o in obj_data]
obj_to_img = unpack_var(obj_to_img)
triple_to_img = unpack_var(triple_to_img)
triples_out = []
obj_data_out = [[] for _ in obj_data]
obj_offset = 0
N = obj_to_img.max() + 1
for i in range(N):
o_idxs = (obj_to_img == i).nonzero().view(-1)
t_idxs = (triple_to_img == i).nonzero().view(-1)
cur_triples = triples[t_idxs].clone()
cur_triples[:, 0] -= obj_offset
cur_triples[:, 2] -= obj_offset
triples_out.append(cur_triples)
for j, o_data in enumerate(obj_data):
cur_o_data = None
if o_data is not None:
cur_o_data = o_data[o_idxs]
obj_data_out[j].append(cur_o_data)
obj_offset += o_idxs.size(0)
return triples_out, obj_data_out
|