You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
|
#
|
|
# 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 numpy as np
|
|
from PIL import Image
|
|
import copy
|
|
import cv2
|
|
from ..registry import PIPELINES
|
|
|
|
|
|
@PIPELINES.register()
|
|
class MultiRestrictSize(object):
|
|
def __init__(self,
|
|
min_size=None,
|
|
max_size=800,
|
|
flip=False,
|
|
multi_scale=[1.3]):
|
|
self.min_size = min_size
|
|
self.max_size = max_size
|
|
self.multi_scale = multi_scale
|
|
self.flip = flip
|
|
assert ((min_size is None)) or ((max_size is None))
|
|
|
|
def __call__(self, sample):
|
|
samples = []
|
|
image = sample['current_img']
|
|
h, w = image.shape[:2]
|
|
for scale in self.multi_scale:
|
|
# Fixed range of scales
|
|
sc = None
|
|
# Align short edge
|
|
if not (self.min_size is None):
|
|
if h > w:
|
|
short_edge = w
|
|
else:
|
|
short_edge = h
|
|
if short_edge > self.min_size:
|
|
sc = float(self.min_size) / short_edge
|
|
else:
|
|
if h > w:
|
|
long_edge = h
|
|
else:
|
|
long_edge = w
|
|
if long_edge > self.max_size:
|
|
sc = float(self.max_size) / long_edge
|
|
|
|
if sc is None:
|
|
new_h = h
|
|
new_w = w
|
|
else:
|
|
new_h = sc * h
|
|
new_w = sc * w
|
|
new_h = int(new_h * scale)
|
|
new_w = int(new_w * scale)
|
|
|
|
if (new_h - 1) % 16 != 0:
|
|
new_h = int(np.around((new_h - 1) / 16.) * 16 + 1)
|
|
if (new_w - 1) % 16 != 0:
|
|
new_w = int(np.around((new_w - 1) / 16.) * 16 + 1)
|
|
|
|
if new_h == h and new_w == w:
|
|
samples.append(sample)
|
|
else:
|
|
new_sample = {}
|
|
for elem in sample.keys():
|
|
if 'meta' in elem:
|
|
new_sample[elem] = sample[elem]
|
|
continue
|
|
tmp = sample[elem]
|
|
if 'label' in elem:
|
|
new_sample[elem] = sample[elem]
|
|
continue
|
|
else:
|
|
flagval = cv2.INTER_CUBIC
|
|
tmp = cv2.resize(tmp,
|
|
dsize=(new_w, new_h),
|
|
interpolation=flagval)
|
|
new_sample[elem] = tmp
|
|
samples.append(new_sample)
|
|
|
|
if self.flip:
|
|
now_sample = samples[-1]
|
|
new_sample = {}
|
|
for elem in now_sample.keys():
|
|
if 'meta' in elem:
|
|
new_sample[elem] = now_sample[elem].copy()
|
|
new_sample[elem]['flip'] = True
|
|
continue
|
|
tmp = now_sample[elem]
|
|
tmp = tmp[:, ::-1].copy()
|
|
new_sample[elem] = tmp
|
|
samples.append(new_sample)
|
|
|
|
return samples
|
|
|
|
|
|
@PIPELINES.register()
|
|
class MultiNorm(object):
|
|
def __call__(self, samples):
|
|
for idx in range(len(samples)):
|
|
sample = samples[idx]
|
|
for elem in sample.keys():
|
|
if 'meta' in elem:
|
|
continue
|
|
tmp = sample[elem]
|
|
if tmp is None:
|
|
continue
|
|
|
|
if tmp.ndim == 2:
|
|
tmp = tmp[:, :, np.newaxis]
|
|
else:
|
|
tmp = tmp / 255.
|
|
tmp -= (0.485, 0.456, 0.406)
|
|
tmp /= (0.229, 0.224, 0.225)
|
|
|
|
tmp = tmp.transpose((2, 0, 1))
|
|
samples[idx][elem] = tmp
|
|
|
|
return samples
|