def create_transformer(self):
train_tforms, eval_tforms = [transforms.Resize(self.args.resize)
], [transforms.Resize(self.args.resize)]
if self.args.random_crop:
train_tforms.append(
transforms.RandomResizedCrop(self.args.size, scale=(0.8, 1.2)))
else:
train_tforms.append(transforms.CenterCrop(self.args.size))
eval_tforms.append(transforms.CenterCrop(self.args.size))
if self.args.flip:
train_tforms.append(transforms.RandomFlipLeftRight())
if self.args.random_color:
train_tforms.append(
transforms.RandomColorJitter(self.args.color_jitter,
self.args.color_jitter,
self.args.color_jitter, 0.1))
train_tforms.extend([
transforms.ToTensor(),
transforms.Normalize(self.args.mean, self.args.std)
])
eval_tforms.extend([
transforms.ToTensor(),
transforms.Normalize(self.args.mean, self.args.std)
])
train_tforms = transforms.Compose(train_tforms)
eval_tforms = transforms.Compose(eval_tforms)
return train_tforms, eval_tforms
def _evaluate(self, val_data):
"""Test on validation dataset."""
acc_top1 = mx.metric.Accuracy()
acc_top5 = mx.metric.TopKAccuracy(min(5, self.num_class))
if not isinstance(val_data, (gluon.data.DataLoader, mx.io.MXDataIter)):
if hasattr(val_data, 'to_mxnet'):
val_data = val_data.to_mxnet()
resize = int(
math.ceil(self.input_size / self._cfg.train.crop_ratio))
normalize = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
transform_test = transforms.Compose([
transforms.Resize(resize, keep_ratio=True),
transforms.CenterCrop(self.input_size),
transforms.ToTensor(), normalize
])
val_data = gluon.data.DataLoader(
val_data.transform_first(transform_test),
batch_size=self._cfg.valid.batch_size,
shuffle=False,
last_batch='keep',
num_workers=self._cfg.valid.num_workers)
for _, batch in enumerate(val_data):
data, label = self.batch_fn(batch, self.ctx)
outputs = [
self.net(X.astype(self._cfg.train.dtype, copy=False))
for X in data
]
acc_top1.update(label, outputs)
acc_top5.update(label, outputs)
_, top1 = acc_top1.get()
_, top5 = acc_top5.get()
return top1, top5
def load_img_batch(dataset, batch_size, type):
if type == 'train':
transform=img_transforms.Compose([\
#随机对图像裁剪出面积为原图像面积的0.08~1倍
#高/宽：3/4 ~ 4/3，最后高度与宽度都缩放到224像素
img_transforms.RandomResizedCrop(224,scale=(0.08,1.0),ratio=(3.0/4.0,4.0/3.0)),\
#随机左右翻转
img_transforms.RandomFlipLeftRight(),\
#随机变化亮度、对比度、饱和度
img_transforms.RandomColorJitter(brightness=0.4,contrast=0.4,saturation=0.4),\
#随机噪声
img_transforms.RandomLighting(0.1),\
img_transforms.ToTensor(),\
# 对图像的每个通道做标准化
img_transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
return gdata.DataLoader(dataset.transform_first(transform),
batch_size=batch_size,
shuffle=True,
last_batch='keep')
elif type == 'test':
transform=img_transforms.Compose([\
img_transforms.Resize(256),\
img_transforms.CenterCrop(224),\
img_transforms.ToTensor(),\
img_transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
return gdata.DataLoader(dataset.transform_first(transform),
batch_size=batch_size,
shuffle=False,
last_batch='keep')
def get_data_loader(data_dir, batch_size, num_workers):
normalize = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
input_size = opt.input_size
crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
resize = int(math.ceil(input_size / crop_ratio))
def batch_fn(batch, ctx):
data = gluon.utils.split_and_load(batch[0],
ctx_list=ctx,
batch_axis=0)
label = gluon.utils.split_and_load(batch[1],
ctx_list=ctx,
batch_axis=0)
return data, label
transform_test = transforms.Compose([
transforms.Resize(resize, keep_ratio=True),
transforms.CenterCrop(input_size),
transforms.ToTensor(), normalize
])
val_data = gluon.data.DataLoader(imagenet.classification.ImageNet(
data_dir, train=False).transform_first(transform_test),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
return val_data, batch_fn
def __init__(self,
datasetroot=None,
istrain=False,
resize_size=None,
crop_size=None,
interpolation=INTERPOLATION,
mean=MEAN,
std=STD):
self.datasetroot = datasetroot
self.istrain = istrain
transform_list = [ImageRead()]
if resize_size is not None:
transform_list.append(
transforms.Resize(resize_size, interpolation=interpolation))
if crop_size is not None:
if self.istrain:
transform_list.append(
RandomCrop(crop_size, interpolation=interpolation))
else:
transform_list.append(
transforms.CenterCrop(crop_size,
interpolation=interpolation))
if self.istrain:
transform_list.append(transforms.RandomFlipLeftRight())
transform_list.append(transforms.ToTensor())
transform_list.append(transforms.Normalize(mean, std))
self.transformer = transforms.Compose(transform_list)
def DogDataLoader(opt):
transform_train = T.Compose([
T.Resize(256),
T.CenterCrop(224),
T.RandomFlipLeftRight(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
transform_valid = T.Compose([
T.Resize(224),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_set = vision.ImageFolderDataset(opt.train_dir, flag=1)
valid_set = vision.ImageFolderDataset(opt.valid_dir, flag=1)
loader = gluon.data.DataLoader
train_loader = loader(train_set.transform_first(transform_train),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
last_batch='rollover')
valid_loader = loader(valid_set.transform_first(transform_valid),
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.num_workers,
last_batch='keep')
return train_loader, valid_loader
def _evaluate(self, val_data):
"""Test on validation dataset."""
acc_top1 = mx.metric.Accuracy()
acc_top5 = mx.metric.TopKAccuracy(min(5, self.num_class))
if not isinstance(val_data, (gluon.data.DataLoader, mx.io.ImageRecordIter)):
from ...tasks.dataset import ImageClassificationDataset
if isinstance(val_data, ImageClassificationDataset):
val_data = val_data.to_mxnet()
transform_test = transforms.Compose([
transforms.Resize(resize, keep_ratio=True),
transforms.CenterCrop(self.input_size),
transforms.ToTensor(),
normalize
])
val_data = gluon.data.DataLoader(
val_data.transform_first(transform_test),
batch_size=self._cfg.valid.batch_size, shuffle=False, last_batch='keep',
num_workers=self._cfg.valid.num_workers)
for _, batch in enumerate(val_data):
data, label = self.batch_fn(batch, self.ctx)
outputs = [self.net(X.astype(self._cfg.train.dtype, copy=False)) for X in data]
acc_top1.update(label, outputs)
acc_top5.update(label, outputs)
_, top1 = acc_top1.get()
_, top5 = acc_top5.get()
return top1, top5
def getNdFromFrame(self):
#use cv2 to capture a frame at self.second from self.source to return imagedata [H x W x RGB]
RootSrc = "C:/Users/samta/TurtleCam" + self.source
vidcap = cv2.VideoCapture(RootSrc)
#perhaps a much quicker method for grabbing frames https://gist.github.com/kylemcdonald/85d70bf53e207bab3775
#This function takes an incredibly long time to execute
vidcap.set(cv2.CAP_PROP_POS_MSEC, self.second * 1000)
(hasFrames, image) = vidcap.read() #stored as 3-D list [H][W][RGB]
if hasFrames:
#reorganize later.. make class in dataAnalyze
img_nump = np.asarray(image)
temp_nd = nd.array(img_nump)
transformer = trans.Compose([
trans.Resize(300),
trans.CenterCrop(255),
trans.ToTensor(),
trans.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
norm_nd_img = transformer(temp_nd)
return hasFrames, norm_nd_img
else:
return hasFrames, 'null'
def transform_image(array):
"""
Should transform image by:
1) Resizing the shortest dimension to 224. e.g (448, 1792) -> (224, 896).
2) Cropping to a center square of dimension (224, 224).
3) Converting the image from HWC layout to CHW layout.
4) Normalizing the image using ImageNet statistics (i.e. per colour channel mean and variance).
5) Creating a batch of 1 image.
:param filepath: array (in HWC layout).
:type filepath: mx.nd.NDArray
:return: a batch of a single transformed images (in NCHW layout)
:rtype: mx.nd.NDArray
"""
# YOUR CODE HERE
train_trans = transforms.Compose([
transforms.Resize(224, keep_ratio=True),
transforms.CenterCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([.485, .456, .606], [.229, .224, .225])
])
return (train_trans(array).expand_dims(0))
raise NotImplementedError()
def get_test_transform(resize, crop, mean, std):
return transforms.Compose([
# transforms.Resize(resize),
transforms.CenterCrop(crop),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
def __init__(self,
name=None,
train_path=None,
val_path=None,
input_size=224,
crop_ratio=0.875,
jitter_param=0.4,
**kwargs):
self.name = name
self.train_path = train_path
self.val_path = val_path
self.input_size = input_size
resize = int(math.ceil(input_size / crop_ratio))
self.transform_train = transforms.Compose([
transforms.Resize(resize),
transforms.RandomResizedCrop(input_size),
transforms.RandomFlipLeftRight(),
transforms.RandomColorJitter(brightness=jitter_param,
contrast=jitter_param,
saturation=jitter_param),
transforms.RandomLighting(0.1),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
self.transform_val = transforms.Compose([
transforms.Resize(resize),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
self._read_dataset(**kwargs)
def mxnet_cifar10(im):
img = image.imread(im)
# plt.imshow(img.asnumpy())
# plt.show()
# transform image
transform_fn = transforms.Compose([
transforms.Resize(32),
transforms.CenterCrop(32),
transforms.ToTensor(),
transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])
])
img = transform_fn(img)
# plt.imshow(nd.transpose(img, (1,2,0)).asnumpy())
# plt.show()
# load pre-trained model
net = get_model('cifar_resnet110_v1', classes=10, pretrained=True)
# predict class
pred = net(img.expand_dims(axis=0))
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck']
ind = nd.argmax(pred, axis=1).astype('int')
return [str(class_names[ind.asscalar()]), str(round(nd.softmax(pred)[0][ind].asscalar(), 2))]
def imagenet_val_transform(ds_metainfo,
mean_rgb=(0.485, 0.456, 0.406),
std_rgb=(0.229, 0.224, 0.225)):
"""
Create image transform sequence for validation subset.
Parameters:
----------
ds_metainfo : DatasetMetaInfo
ImageNet-1K dataset metainfo.
mean_rgb : tuple of 3 float
Mean of RGB channels in the dataset.
std_rgb : tuple of 3 float
STD of RGB channels in the dataset.
Returns
-------
Sequential
Image transform sequence.
"""
input_image_size = ds_metainfo.input_image_size
resize_value = calc_val_resize_value(
input_image_size=ds_metainfo.input_image_size,
resize_inv_factor=ds_metainfo.resize_inv_factor)
return transforms.Compose([
transforms.Resize(
size=resize_value,
keep_ratio=True),
transforms.CenterCrop(size=input_image_size),
transforms.ToTensor(),
transforms.Normalize(
mean=mean_rgb,
std=std_rgb)
])
def get_data_loader(data_dir, batch_size, num_workers):
normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
jitter_param = 0.4
lighting_param = 0.1
input_size = opt.input_size
def batch_fn(batch, ctx):
data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)
return data, label
transform_train = transforms.Compose([
transforms.RandomResizedCrop(input_size),
transforms.RandomFlipLeftRight(),
transforms.RandomColorJitter(brightness=jitter_param, contrast=jitter_param,
saturation=jitter_param),
transforms.RandomLighting(lighting_param),
transforms.ToTensor(),
normalize
])
transform_test = transforms.Compose([
transforms.Resize(256, keep_ratio=True),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
normalize
])
train_data = gluon.data.DataLoader(
imagenet.classification.ImageNet(data_dir, train=True).transform_first(transform_train),
batch_size=batch_size, shuffle=True, last_batch='discard', num_workers=num_workers)
val_data = gluon.data.DataLoader(
imagenet.classification.ImageNet(data_dir, train=False).transform_first(transform_test),
batch_size=batch_size, shuffle=False, num_workers=num_workers)
return train_data, val_data, batch_fn
def imagenet_val_transform(ds_metainfo):
"""
Create image transform sequence for validation subset.
Parameters:
----------
ds_metainfo : DatasetMetaInfo
ImageNet-1K dataset metainfo.
Returns
-------
Sequential
Image transform sequence.
"""
input_image_size = ds_metainfo.input_image_size
resize_value = calc_val_resize_value(
input_image_size=ds_metainfo.input_image_size,
resize_inv_factor=ds_metainfo.resize_inv_factor)
return transforms.Compose([
transforms.Resize(size=resize_value,
keep_ratio=True,
interpolation=ds_metainfo.interpolation),
transforms.CenterCrop(size=input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=ds_metainfo.mean_rgb,
std=ds_metainfo.std_rgb)
])
def preprocess_test_data(normalize):
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
])
return transform_test
def fdv1_val_transform(ds_metainfo,
mean_rgb=(0.485, 0.456, 0.406),
std_rgb=(0.229, 0.224, 0.225)):
"""
Create image transform sequence for validation subset.
Parameters:
----------
ds_metainfo : DatasetMetaInfo
FDV1 dataset metainfo.
mean_rgb : tuple of 3 float
Mean of RGB channels in the dataset.
std_rgb : tuple of 3 float
STD of RGB channels in the dataset.
Returns
-------
Sequential
Image transform sequence.
"""
input_image_size = ds_metainfo.input_image_size
resize_value = calc_val_resize_value(
input_image_size=ds_metainfo.input_image_size,
resize_inv_factor=ds_metainfo.resize_inv_factor)
if ds_metainfo.do_downscale:
transform_list = [DownscaleOnDemand(size=150, threshold_size=256)]
else:
transform_list = []
transform_list += [
Squaring(size=resize_value),
transforms.CenterCrop(size=input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=mean_rgb, std=std_rgb)
]
return transforms.Compose(transform_list)
def test_model_for_ml(self):
net_path = os.path.join(
DATA_DIR, 'model',
'epoch-3-0.48-20180920164709.params-symbol.json')
params_path = os.path.join(
DATA_DIR, 'model',
'epoch-3-0.48-20180920164709.params-0003.params')
net = gluon.nn.SymbolBlock.imports(net_path, ['data'], params_path)
im_path = os.path.join(DATA_DIR, 'imgs_data',
'd4YE10xHdvbwKJV5yBYsoJJke6K9b.jpg')
img = image.imread(im_path)
# plt.imshow(img.asnumpy())
# plt.show()
transform_fn = transforms.Compose([
transforms.Resize(224, keep_ratio=True),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
img = transform_fn(img)
img = nd.expand_dims(img, axis=0)
res = net(img.as_in_context(self.ctx[0]))
res = sigmoid(nd.squeeze(res)).asnumpy()
res = np.where(res > 0.5, 1, 0)
indexes, = np.where(res == 1)
res_classes = [self.class_names[i] for i in indexes.tolist()]
# print(indexes.tolist())
print('测试类别: {}'.format(res_classes))
def get_val_data_source(dataset_dir,
batch_size,
num_workers,
input_image_size=(224, 224),
resize_inv_factor=0.875):
assert (resize_inv_factor > 0.0)
if isinstance(input_image_size, int):
input_image_size = (input_image_size, input_image_size)
resize_value = int(
math.ceil(float(input_image_size[0]) / resize_inv_factor))
mean_rgb = (0.485, 0.456, 0.406)
std_rgb = (0.229, 0.224, 0.225)
transform_val = transforms.Compose([
transforms.Resize(resize_value, keep_ratio=True),
transforms.CenterCrop(input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=mean_rgb, std=std_rgb)
])
return gluon.data.DataLoader(dataset=CUB200_2011(
root=dataset_dir, train=False).transform_first(fn=transform_val),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
def predict(self, X, input_size=224, plot=True):
""" This method should be able to produce predictions regardless if:
X = single data example (e.g. single image, single document),
X = batch of many examples, X = task.Dataset object
"""
"""The task predict function given an input.
Args:
img: the input
Example:
>>> ind, prob = classifier.predict('example.jpg')
"""
# load and display the image
img = mx.image.imread(X) if isinstance(
X, str) and os.path.isfile(X) else X
if plot:
plt.imshow(img.asnumpy())
plt.show()
# model inference
input_size = self.model.input_size if hasattr(
self.model, 'input_size') else input_size
resize = int(math.ceil(input_size / 0.875))
transform_fn = transforms.Compose([
transforms.Resize(resize),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
img = transform_fn(img)
proba = self.predict_proba(img)
ind = mx.nd.argmax(proba, axis=1).astype('int')
idx = mx.nd.stack(mx.nd.arange(proba.shape[0], ctx=proba.context),
ind.astype('float32'))
return ind, mx.nd.gather_nd(proba, idx)
def get_data_loader(data_dir, batch_size, num_workers):
normalize = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
def batch_fn(batch, ctx):
data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
label = gluon.utils.split_and_load(batch[1],
ctx_list=ctx,
batch_axis=0)
return data, label
if opt.mode == 'symbolic':
val_data = mx.io.NDArrayIter(
mx.nd.random.normal(shape=(opt.dataset_size, 3, 224, 224),
ctx=context),
label=mx.nd.array(range(opt.dataset_size)),
batch_size=batch_size,
)
transform_test = transforms.Compose([
transforms.Resize(256, keep_ratio=True),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
])
val_data = gluon.data.DataLoader(imagenet.classification.ImageNet(
data_dir, train=False).transform_first(transform_test),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
return val_data, batch_fn
def create_loader(self):
"""
Create the data loader
:return:
"""
if self.args.mode.upper() == 'TRAIN':
tforms = []
tforms.append(transforms.Resize(self.args.resize))
if self.args.flip:
tforms.append(transforms.RandomFlipLeftRight())
if self.args.random_crop:
tforms.append(
transforms.RandomResizedCrop(self.args.im_size,
scale=(0.8, 1)))
else:
tforms.append(transforms.CenterCrop(self.args.im_size))
if self.args.random_jitter:
tforms.append(transforms.RandomColorJitter(0.4, 0.4, 0.4, 0.4))
tforms.append(transforms.ToTensor())
tforms.append(
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225)))
tforms = transforms.Compose(tforms)
tr_db = list(self.cfg['train'].values())[0]
dataset = ImageRecordDataset(tr_db['rec'], transform=tforms)
self.tr_loader = DataLoader(dataset,
batch_size=self.args.bs,
num_workers=8,
pin_memory=True)
else:
tforms = transforms.Compose([
transforms.Resize(self.args.resize),
transforms.CenterCrop(self.args.im_size),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))
])
self.eval_tforms = tforms
def load_images(images_file_path,
batch_size,
resize_size=256,
is_train=True,
crop_size=224,
is_cen=False,
sampler=None,
pseudo_labels=None):
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if not is_train:
transformer = transforms.Compose([
transforms.Resize(resize_size),
transforms.CenterCrop(crop_size),
transforms.ToTensor(), normalize
])
shuffle = False
last_bacth = 'keep'
else:
if is_cen:
transformer = transforms.Compose([
transforms.Resize(resize_size),
transforms.CenterCrop(crop_size),
transforms.RandomFlipLeftRight(),
transforms.ToTensor(), normalize
])
else:
transformer = transforms.Compose([
transforms.Resize(resize_size),
transforms.RandomResizedCrop(crop_size, scale=(0.8, 1.0)),
transforms.RandomFlipLeftRight(),
transforms.ToTensor(), normalize
])
shuffle = False if sampler is not None else True
last_bacth = 'keep'
imageset = ImageFolderDataset(images_file_path,
pseudo_labels=pseudo_labels)
data_loader = DataLoader(dataset=imageset.transform_first(transformer),
shuffle=shuffle,
batch_size=batch_size,
last_batch=last_bacth,
sampler=sampler,
num_workers=0)
return data_loader
def transform_cifar10_dataset_val():
"""
Should create a transformation that performs center cropping
:return: A gluon transform object
:rtype: gluon.Block
"""
transforms_val = transforms.CenterCrop(16)
return transforms_val
def transform_image(self, img):
transform_fn = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
img = transform_fn(mx.nd.array(img)).asnumpy()
return img
def transform(img):
input_size = (64, 64)
h = len(img)
w = len(img[0])
transform = transforms.Compose([
transforms.CenterCrop(w,w),
transforms.Resize(input_size, keep_ratio = True),
transforms.ToTensor()
])
return transform(img)
def transform_image(img):
transform_fn = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
img = transform_fn(mx.nd.array(img)).asnumpy()
img = np.expand_dims(img, axis=0) # batchify
return img
def preprocess(self, img):
transform_fn = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
img = transform_fn(img)
img = img.expand_dims(axis=0)
return img
def imagenet_val_transform(input_image_size=(224, 224),
mean_rgb=(0.485, 0.456, 0.406),
std_rgb=(0.229, 0.224, 0.225),
resize_value=256):
return transforms.Compose([
transforms.Resize(size=resize_value, keep_ratio=True),
transforms.CenterCrop(size=input_image_size),
transforms.ToTensor(),
transforms.Normalize(mean=mean_rgb, std=std_rgb)
])
def config_resnet_matcher(all_imgs,
cache_dir: Path,
conf_thresh: float = 0.92557,
ctx=mx.gpu(0)):
mod = gluoncv.model_zoo.get_model('ResNet50_v2', ctx=ctx, pretrained=True)
mod.fc = gluon.contrib.nn.Identity()
mod.hybridize()
transform_fn = transforms.Compose([
transforms.Resize(256, keep_ratio=True),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
embs_path = ensure_path(cache_dir / 'embeddings')
mapping_path = cache_dir / 'mapping.txt'
embs = []
path2emb = {}
if mapping_path.exists():
loaded = []
with open(str(mapping_path), 'r') as f:
for line in f:
parts = line.strip().split(';')
left = Path(parts[0])
loaded.append(left)
path2emb[left] = np.load(str(embs_path /
parts[1]))['embedding']
all_imgs = list(set(all_imgs) - set(loaded))
if len(all_imgs) > 0:
dataset = Loader(all_imgs)
loaded = gluon.data.DataLoader(dataset.transform_first(transform_fn),
batch_size=64,
shuffle=False,
num_workers=3,
pin_memory=True)
print('Preparing matcher')
for batch in tqdm(loaded, total=len(loaded)):
inp = batch.as_in_context(mx.gpu(0))
pred = mod(inp)
embs.append(pred.asnumpy())
embs = np.concatenate(embs, axis=0)
normed = embs / np.sqrt(np.sum(embs**2, axis=-1, keepdims=True))
for cur_path, cur_emb in zip(all_imgs, normed):
path2emb[cur_path] = cur_emb
with open(str(mapping_path), 'a') as f:
print(f'{cur_path};{cur_path.stem}.npz', file=f)
np.savez(str(embs_path / f'{cur_path.stem}.npz'),
embedding=cur_emb)
def match(left: Path, right: Path):
dist = np.dot(path2emb[left], path2emb[right])
return dist > conf_thresh
return match
