class FeatureExtractionDemo(object):
def __init__(self, cfg, parallel=False):
"""
Args:
cfg (CfgNode):
parallel (bool) whether to run the model in different processes from visualization.:
Useful since the visualization logic can be slow.
"""
self.cfg = cfg
self.parallel = parallel
if parallel:
self.num_gpus = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, self.num_gpus)
else:
self.predictor = DefaultPredictor(cfg)
def run_on_image(self, original_image):
"""
Args:
original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
This is the format used by OpenCV.
Returns:
predictions (np.ndarray): normalized feature of the model.
"""
# the model expects RGB inputs
original_image = original_image[:, :, ::-1]
# Apply pre-processing to image.
image = cv2.resize(original_image, tuple(self.cfg.INPUT.SIZE_TEST[::-1]), interpolation=cv2.INTER_CUBIC)
# Make shape with a new batch dimension which is adapted for
# network input
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))[None]
predictions = self.predictor(image)
return predictions
def run_on_loader(self, data_loader):
if self.parallel:
buffer_size = self.predictor.default_buffer_size
batch_data = deque()
for cnt, batch in enumerate(data_loader):
batch_data.append(batch)
self.predictor.put(batch["images"])
if cnt >= buffer_size:
batch = batch_data.popleft()
predictions = self.predictor.get()
yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
while len(batch_data):
batch = batch_data.popleft()
predictions = self.predictor.get()
yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
else:
for batch in data_loader:
predictions = self.predictor(batch["images"])
yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
def __init__(self, cfg, parallel=False):
"""
Args:
cfg (CfgNode):
parallel (bool) whether to run the model in different processes from visualization.:
Useful since the visualization logic can be slow.
"""
self.cfg = cfg
self.parallel = parallel
if parallel:
self.num_gpus = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, self.num_gpus)
else:
self.predictor = DefaultPredictor(cfg)
def __init__(self, cfgs):
"""
Args:
cfg (CfgNode):
"""
self.cfgs = cfgs
self.predictors = []
for cfg in cfgs:
self.predictors.append(DefaultPredictor(cfg))
def run(self):
predictor = DefaultPredictor(self.cfg)
while True:
task = self.task_queue.get()
if isinstance(task, AsyncPredictor._StopToken):
break
idx, data = task
result = predictor(data)
self.result_queue.put((idx, result))
def __init__(self, cfg, device='cuda:0', parallel=False):
"""
Args:
cfg (CfgNode):
parallel (bool) whether to run the model in different processes from visualization.:
Useful since the visualization logic can be slow.
"""
self.cfg = cfg
self.parallel = parallel
if parallel:
self.num_gpus = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, self.num_gpus)
else:
self.predictor = DefaultPredictor(cfg, device)
num_channels = len(cfg.MODEL.PIXEL_MEAN)
self.mean = torch.tensor(cfg.MODEL.PIXEL_MEAN).view(
1, num_channels, 1, 1)
self.std = torch.tensor(cfg.MODEL.PIXEL_STD).view(
1, num_channels, 1, 1)
class FeatureExtractionDemo(object):
def __init__(self, cfg, device='cuda:0', parallel=False):
"""
Args:
cfg (CfgNode):
parallel (bool) whether to run the model in different processes from visualization.:
Useful since the visualization logic can be slow.
"""
self.cfg = cfg
self.parallel = parallel
if parallel:
self.num_gpus = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, self.num_gpus)
else:
self.predictor = DefaultPredictor(cfg, device)
num_channels = len(cfg.MODEL.PIXEL_MEAN)
self.mean = torch.tensor(cfg.MODEL.PIXEL_MEAN).view(
1, num_channels, 1, 1)
self.std = torch.tensor(cfg.MODEL.PIXEL_STD).view(
1, num_channels, 1, 1)
def run_on_image(self, original_image):
"""
Args:
original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
This is the format used by OpenCV.
Returns:
predictions (np.ndarray): normalized feature of the model.
"""
# the model expects RGB inputs
original_image = original_image[:, :, ::-1]
# Apply pre-processing to image.
image = cv2.resize(original_image,
tuple(self.cfg.INPUT.SIZE_TEST[::-1]),
interpolation=cv2.INTER_CUBIC)
image = torch.as_tensor(image.astype("float32").transpose(2, 0,
1))[None]
image.sub_(self.mean).div_(self.std)
predictions = self.predictor(image)
return predictions
def run_on_loader(self, data_loader):
image_gen = self._image_from_loader(data_loader)
if self.parallel:
buffer_size = self.predictor.default_buffer_size
batch_data = deque()
for cnt, batch in enumerate(image_gen):
batch_data.append(batch)
self.predictor.put(batch['images'])
if cnt >= buffer_size:
batch = batch_data.popleft()
predictions = self.predictor.get()
yield predictions, batch['targets'].numpy(
), batch['camid'].numpy()
while len(batch_data):
batch = batch_data.popleft()
predictions = self.predictor.get()
yield predictions, batch['targets'].numpy(
), batch['camid'].numpy()
else:
for batch in image_gen:
predictions = self.predictor(batch['images'])
yield predictions, batch['targets'].numpy(
), batch['camid'].numpy()
def _image_from_loader(self, data_loader):
data_loader.reset()
data = data_loader.next()
while data is not None:
yield data
data = data_loader.next()
)
parser.add_argument(
"--opts",
help="Modify config options using the command-line 'KEY VALUE' pairs",
default=[],
nargs=argparse.REMAINDER,
)
return parser
if __name__ == '__main__':
args = get_parser().parse_args()
cfg = setup_cfg(args)
demo = DefaultPredictor(cfg)
feats = []
if args.input:
if len(args.input) == 1:
args.input = glob.glob(os.path.expanduser(args.input[0]))
assert args.input, "The input path(s) was not found"
for path in tqdm.tqdm(args.input, disable=not args.output):
img = cv2.imread(path)
feats.append(demo(img))
cos_12 = np.dot(feats[0], feats[1].T).item()
cos_13 = np.dot(feats[0], feats[2].T).item()
cos_23 = np.dot(feats[1], feats[2].T).item()
print('cosine similarity is {:.4f}, {:.4f}, {:.4f}'.format(cos_12, cos_13, cos_23))
