xautodl/lib/utils/fb_transform.py
2019-02-01 01:27:38 +11:00

15 lines
411 B
Python

import torch
import random
import numpy as np
class ApplyOffset(object):
def __init__(self, offset):
assert isinstance(offset, int), 'The offset is not right : {}'.format(offset)
self.offset = offset
def __call__(self, x):
if isinstance(x, np.ndarray) and x.dtype == 'uint8':
x = x.astype(int)
if isinstance(x, np.ndarray) and x.size == 1:
x = int(x)
return x + self.offset