def __init__(self, args, train):
self.path = args.datapath
self.videoshape = args.videoshape
self.train = train
self.dataset = args.dataset
if self.train:
self.trainlist = './datasets/' + str(
args.dataset) + '_%d' % (args.cropshape[1]) + '.txt'
self.imglist = default_list_reader(self.trainlist)
self.train_dir = join(self.path, args.dataset, 'training')
self.train_transform = transforms.Compose([
ToFloatTensor3D(),
Normalize(),
ToSpatialCrops(args.videoshape, args.cropshape)
])
else:
self.test_dir = join(self.path, args.dataset, 'testing')
self.cur_video_id = None
self.cur_video_frames = None
self.cur_video_gt = None
self.cur_background = None
self.cur_len = 0
self.test_transform = transforms.Compose([
RemoveBackground(threshold=128),
ToFloatTensor3D(),
Normalize(),
ToCenterCrops(args.videoshape, args.cropshape)
])
def regression(self, dpair):
model = self.model
rescale = Rescale(MODEL_CONFIG['input_size'])
try:
prev_im, curr_im = rescale({'image':
dpair['prev_im']})['image'], rescale(
{'image':
dpair['curr_im']})['image']
except:
return np.array([0, 0, 0, 0])
trans_tomodel = transforms.Compose([Normalize(), ToTensor()])
dpair = trans_tomodel({'prev_im': prev_im, 'curr_im': curr_im})
prev_im, curr_im = dpair['prev_im'], dpair['curr_im']
if self.use_gpu is True:
prev_im, curr_im = Variable(prev_im.cuda()), Variable(
curr_im.cuda())
model = model.cuda()
else:
prev_im, curr_im = Variable(prev_im), Variable(curr_im)
prev_im = prev_im[None, :, :, :]
curr_im = curr_im[None, :, :, :]
regression_bbox = model(prev_im, curr_im)
bbox = regression_bbox.data.cpu().numpy()
bbox = bbox[0, :] / MODEL_CONFIG['bbox_scale']
return bbox
def __init__(self, mode, model=None):
"""TODO: to be defined.
:mode: TODO
:model: TODO
"""
assert isinstance(mode, str)
assert mode in pconfig.AVAILABLE_MODES
if mode in ['fully-supervised', 'val', 'test']:
assert model is None
elif mode == 'weakly-supervised':
assert model is not None
assert isinstance(model, CRAFT)
else:
raise ValueError
self._mode = mode
# NOTE [ Data Transformation ]
# 1. Must-have steps:
# 1.1. Heatmap generation (region_map, affinity_map, confidence_map)
# ----- NOTE [ Data Augmentation ] -----
# 2. A permutation of color variations:
# - Brightness
# - Contrast
# - Saturation
# - Hue
# - ...
# 3. A permutation of geometry transformations:
# - Horizontal flipping
# - Cropping
# - Scaling (input image size reduces)
# - Padding (input image size doesn't change)
# - Rotation
# - Translation
# - Shearing
# - Resizing (variable input)
# ----- END [ Data Augmentation ] -----
# 1.2. ToTensor
# 1.3. Normalization (mean-std, ...)
if self._mode in ['fully-supervised', 'weakly-supervised', 'val']:
self._generate_heatmap = GenerateHeatMap(self._mode, model)
self._resize = Resize(pconfig.TRAIN_IMAGE_SIZE)
elif self._mode == 'test':
self._resize = Resize(pconfig.TEST_IMAGE_SIZE)
else:
raise ValueError
self._to_tensor = ToTensor()
self._normalize =\
Normalize(mean=dconfig.MEAN,
std=dconfig.STD,
inplace=False)
def setup(self):
train_ops = []
valid_ops = []
# Image cropping for the train set
if self.hparams.random_crop:
train_ops.append(MaskedRandomCrop(self.hparams.train_inp_size))
else:
train_ops.append(CenterCrop(self.hparams.train_inp_size))
# Image cropping for the valid set
valid_ops.append(CentroidCrop(self.hparams.valid_inp_size))
# ToTensor
train_ops.append(ToTensor())
valid_ops.append(ToTensor())
# Normalize
train_ops.append(Normalize(mean=self.mean, std=self.std))
valid_ops.append(Normalize(mean=self.mean, std=self.std))
# Data split into train, valid and test sets.
indices = list(range(len(self.data)))
test_dataset = Subset(self.data, indices[:self.hparams.test_size])
valid_dataset = Subset(
self.data, indices[self.hparams.test_size:self.hparams.test_size +
self.hparams.valid_size])
train_dataset = Subset(
self.data,
indices[self.hparams.test_size + self.hparams.valid_size:])
self.datasets['train'] = RealHandsDataset(train_dataset,
Compose(train_ops))
self.datasets['valid'] = RealHandsDataset(valid_dataset,
Compose(valid_ops))
self.datasets['test'] = RealHandsDataset(test_dataset,
Compose(valid_ops))
self.datasets['mean_std'] = RealHandsDataset(self.data, ToTensor())
def preprocess(self, curr_search_reg, target_pad, curr_bb_recenter):
assert(len(curr_search_reg) == len(target_pad) == len(curr_bb_recenter)), "Length not legal..."
minibs = len(curr_search_reg)
def bboxscale(image, bbox):
h, w = image.shape[0], image.shape[1]
bbox = np.array([bbox[0]/w, bbox[1]/h, bbox[2]/w, bbox[3]/h])
bbox *= MODEL_CONFIG['bbox_scale']
return torch.from_numpy(bbox)
trans = Compose([Rescale(MODEL_CONFIG['input_size']),
Normalize(),
ToTensor()])
for i in range(minibs):
self.prev_im.append(trans(target_pad[i]))
self.curr_im.append(trans(curr_search_reg[i]))
self.curr_bb.append(bboxscale(curr_search_reg[i], curr_bb_recenter[i]))