def main():
# 1. Get input arguments
args = get_args()
# 2. Create config instance from args above
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
# 3. Create DataManager Instance
datamanager = build_datamanager(cfg)
print('Building model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(
name=cfg.model.name,
num_classes=datamanager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
if cfg.use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = torchreid.optim.build_lr_scheduler(optimizer,
**lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer,
scheduler=scheduler)
print('Building {}-engine for {}-reid'.format(cfg.loss.name,
cfg.data.type))
# Build engine and run
engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(cfg))
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='path to config file')
parser.add_argument('--output_name', type=str, default='model')
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
model = torchreid.models.build_model(name=cfg.model.name,
num_classes=2,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu,
dropout_prob=cfg.model.dropout_prob,
feature_dim=cfg.model.feature_dim,
activation=cfg.model.activation,
in_first=cfg.model.in_first)
load_pretrained_weights(model, cfg.model.load_weights)
model.eval()
_, transform = build_transforms(cfg.data.height, cfg.data.width)
input_size = (cfg.data.height, cfg.data.width, 3)
img = np.random.rand(*input_size).astype(np.float32)
img = np.uint8(img * 255)
im = Image.fromarray(img)
blob = transform(im).unsqueeze(0)
torch.onnx.export(model,
blob,
args.output_name + '.onnx',
verbose=True,
export_params=True)
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='path to config file')
parser.add_argument('-s',
'--sources',
type=str,
nargs='+',
help='source datasets (delimited by space)')
parser.add_argument('-t',
'--targets',
type=str,
nargs='+',
help='target datasets (delimited by space)')
parser.add_argument('--transforms',
type=str,
nargs='+',
help='data augmentation')
parser.add_argument('--root',
type=str,
default='',
help='path to data root')
parser.add_argument(
'--gpu-devices',
type=str,
default='',
)
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
if cfg.use_gpu and args.gpu_devices:
# if gpu_devices is not specified, all available gpus will be used
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
datamanager = build_datamanager(cfg)
print('Building model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(
name=cfg.model.name,
num_classes=datamanager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
if cfg.use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = torchreid.optim.build_lr_scheduler(optimizer,
**lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer)
print('Building {}-engine for {}-reid'.format(cfg.loss.name,
cfg.data.type))
engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(cfg))
else:
ret.append(i)
return ret
if __name__ == '__main__':
# __author__ == '__toka__'
# d = MOT16(root='/mnt/lustre/share/fengweitao')
d = CrowdHuman(
root='/mnt/lustre/share/fengweitao',
meta_file=
'/mnt/lustre/share/fengweitao/crowd_human/annotation_train.odgt')
dl = torch.utils.data.DataLoader(d, batch_size=16, num_workers=4)
print(dl)
config_file = 'configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml'
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
cfg.merge_from_file(config_file)
model = torchreid.models.build_model(name=cfg.model.name,
num_classes=1024,
loss=cfg.loss.name,
pretrained=False,
use_gpu=cfg.use_gpu)
load_pretrained_weights(model, cfg.model.load_weights)
# m = None
model.eval()
model.cuda()
all_results = []
for i, data in enumerate(dl):
data = to_cuda(data, device='cuda')[0]
print(data.keys())
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='path to config file')
parser.add_argument(
'--gpu-devices',
type=str,
default='',
)
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
if args.config_file:
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
torch.backends.cudnn.benchmark = True
datamanager = ImageDataManager(**imagedata_kwargs(cfg))
trainloader, queryloader, galleryloader = datamanager.return_dataloaders()
print('Building model: {}'.format(cfg.model.name))
model = build_model(cfg.model.name,
datamanager.num_train_pids,
'softmax',
pretrained=cfg.model.pretrained)
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
model = nn.DataParallel(model).cuda()
criterion = CrossEntropyLoss(datamanager.num_train_pids,
label_smooth=cfg.loss.softmax.label_smooth)
optimizer = build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = build_lr_scheduler(optimizer, **lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer)
if cfg.test.evaluate:
distmat = evaluate(model,
queryloader,
galleryloader,
dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature,
rerank=cfg.test.rerank,
return_distmat=True)
if cfg.test.visrank:
visualize_ranked_results(distmat,
datamanager.return_testdataset(),
'image',
width=cfg.data.width,
height=cfg.data.height,
save_dir=osp.join(cfg.data.save_dir,
'visrank'))
return
time_start = time.time()
print('=> Start training')
for epoch in range(cfg.train.start_epoch, cfg.train.max_epoch):
train(epoch,
cfg.train.max_epoch,
model,
criterion,
optimizer,
trainloader,
fixbase_epoch=cfg.train.fixbase_epoch,
open_layers=cfg.train.open_layers)
scheduler.step()
if (epoch + 1) % cfg.test.eval_freq == 0 or (epoch +
1) == cfg.train.max_epoch:
rank1 = evaluate(model,
queryloader,
galleryloader,
dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature,
rerank=cfg.test.rerank)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'rank1': rank1,
'optimizer': optimizer.state_dict(),
}, cfg.data.save_dir)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
import os
import struct
import collections
import redis
import numpy as np
from default_config import get_default_config
from utils import get_features_one, get_dist_one
# ReID
cfg_reid = get_default_config()
body_bank_test = {}
# Redis
r = redis.Redis(
host=os.getenv('SMART_SPACE_TRACKING_REDIS_PORT_6379_TCP_ADDR'),
port=6379,
db=0)
pipe = r.pipeline()
p = r.pubsub()
def reid(color_boxes_of_frame, track_bbs_ids_for_reid):
track_ids_for_reid = [int(b[-1]) for b in track_bbs_ids_for_reid]
result_arr_id = np.array([-1] * len(color_boxes_of_frame)).astype(np.int32)
if len(body_bank_test) == 0:
for index in range(len(color_boxes_of_frame)):
# extract features
features_img_query = get_features_one(color_boxes_of_frame[index])
def load_model(model_dir, cnn):
cnn.load_state_dict(torch.load(os.path.join(model_dir, 'cnn.pkl')))
if __name__ == '__main__':
###############################################
# Preparation #
###############################################
# Configuration priority 0 > 1 > 2 > 3 > 4 > 5:
# 0. (only for model configs) loaded model config
# 1. command line options
# 2. default command line options
# 3. command line config file
# 4. main config file
# 5. defult
args = default_config.get_default_config()
configuration.update_config(args)
parser = argparse.ArgumentParser(
description='CNN text multi-class classificer')
# options
parser.add_argument(
'--config',
type=str,
metavar='CONFIG',
help=
'use this configuration file instead of the default config, like "configs.config"'
)
parser.add_argument('--test',
action='store_true',
default=False,
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='path to config file')
parser.add_argument('-s',
'--sources',
type=str,
nargs='+',
help='source datasets (delimited by space)')
parser.add_argument('-t',
'--targets',
type=str,
nargs='+',
help='target datasets (delimited by space)')
parser.add_argument('--transforms',
type=str,
nargs='+',
help='data augmentation')
parser.add_argument('--root',
type=str,
default='',
help='path to data root')
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
datamanager = build_datamanager(cfg)
model_factory = {
'resnet50_fc512': resnet50_fc512,
'osnet_x1_0': osnet_x1_0,
# mixstyle models
'resnet50_fc512_ms12_a0d1': resnet50_fc512_ms12_a0d1,
'resnet50_fc512_ms12_a0d2': resnet50_fc512_ms12_a0d2,
'resnet50_fc512_ms12_a0d3': resnet50_fc512_ms12_a0d3,
'resnet50_fc512_ms12_a0d1_domprior': resnet50_fc512_ms12_a0d1_domprior,
'osnet_x1_0_ms23_a0d1': osnet_x1_0_ms23_a0d1,
'osnet_x1_0_ms23_a0d2': osnet_x1_0_ms23_a0d2,
'osnet_x1_0_ms23_a0d3': osnet_x1_0_ms23_a0d3,
'osnet_x1_0_ms23_a0d1_domprior': osnet_x1_0_ms23_a0d1_domprior,
# ablation
'resnet50_fc512_ms1_a0d1': resnet50_fc512_ms1_a0d1,
'resnet50_fc512_ms123_a0d1': resnet50_fc512_ms123_a0d1,
'resnet50_fc512_ms1234_a0d1': resnet50_fc512_ms1234_a0d1,
'resnet50_fc512_ms14_a0d1': resnet50_fc512_ms14_a0d1,
'resnet50_fc512_ms23_a0d1': resnet50_fc512_ms23_a0d1,
# dropblock models
'resnet50_fc512_db12': resnet50_fc512_db12,
'osnet_x1_0_db23': osnet_x1_0_db23
}
print('Building model: {}'.format(cfg.model.name))
model = model_factory[cfg.model.name](
num_classes=datamanager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
if cfg.use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = torchreid.optim.build_lr_scheduler(optimizer,
**lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer,
scheduler=scheduler)
print('Building {}-engine for {}-reid'.format(cfg.loss.name,
cfg.data.type))
engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(cfg))
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='path to config file')
parser.add_argument('-s',
'--sources',
type=str,
nargs='+',
help='source datasets (delimited by space)')
parser.add_argument('-t',
'--targets',
type=str,
nargs='+',
help='target datasets (delimited by space)')
parser.add_argument('--transforms',
type=str,
nargs='+',
help='data augmentation')
parser.add_argument('--root',
type=str,
default='',
help='path to data root')
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))
print('Building model-1: {}'.format(cfg.model.name))
model1 = torchreid.models.build_model(
name=cfg.model.name,
num_classes=datamanager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
num_params, flops = compute_model_complexity(
model1, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
print('Copying model-1 to model-2')
model2 = copy.deepcopy(model1)
if cfg.model.load_weights1 and check_isfile(cfg.model.load_weights1):
load_pretrained_weights(model1, cfg.model.load_weights1)
if cfg.model.load_weights2 and check_isfile(cfg.model.load_weights2):
load_pretrained_weights(model2, cfg.model.load_weights2)
if cfg.use_gpu:
model1 = nn.DataParallel(model1).cuda()
model2 = nn.DataParallel(model2).cuda()
optimizer1 = torchreid.optim.build_optimizer(model1,
**optimizer_kwargs(cfg))
scheduler1 = torchreid.optim.build_lr_scheduler(optimizer1,
**lr_scheduler_kwargs(cfg))
optimizer2 = torchreid.optim.build_optimizer(model2,
**optimizer_kwargs(cfg))
scheduler2 = torchreid.optim.build_lr_scheduler(optimizer2,
**lr_scheduler_kwargs(cfg))
if cfg.model.resume1 and check_isfile(cfg.model.resume1):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume1,
model1,
optimizer=optimizer1,
scheduler=scheduler1)
if cfg.model.resume2 and check_isfile(cfg.model.resume2):
resume_from_checkpoint(cfg.model.resume2,
model2,
optimizer=optimizer2,
scheduler=scheduler2)
print('Building DML-engine for image-reid')
engine = ImageDMLEngine(datamanager,
model1,
optimizer1,
scheduler1,
model2,
optimizer2,
scheduler2,
margin=cfg.loss.triplet.margin,
weight_t=cfg.loss.triplet.weight_t,
weight_x=cfg.loss.triplet.weight_x,
weight_ml=cfg.loss.dml.weight_ml,
use_gpu=cfg.use_gpu,
label_smooth=cfg.loss.softmax.label_smooth,
deploy=cfg.model.deploy)
engine.run(**engine_run_kwargs(cfg))
def main():
cfg = get_default_config()
# init detector network
print('init detector network...')
odapi = DetectorAPI(path_to_ckpt=cfg.detector.load_weights)
print("done!")
# init intel camera
print('init intel camera...')
pipe = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
# Start streaming
pipe.start(config)
print("done!")
# init banks
print("init gallary...")
body_bank = [] # gallery array of persons
body_bank_dist = [] #
# body_bank_bb = []
body_im_array = []
print("done!")
if cfg.visualisation.key:
print("start visualisation mode...")
plt.ion()
print("done!")
curTime = 0
with tf.Session() as sess:
args_save = True
if args_save:
video_writer = cv2.VideoWriter(cfg.video.path,
cv2.VideoWriter_fourcc(*'XVID'), 6,
(cfg.image.width, cfg.image.height))
while True:
curTime += 1
#print(curTime)
#-----------------------------------------
# pipe to get image from intel cam
frameset = pipe.wait_for_frames()
color_frame = frameset.get_color_frame()
color = np.asanyarray(color_frame.get_data()).copy()
#print(np.shape(color), type(color))
# rgb_frame = make_rgb(color)
frameL = cv2.resize(color, (cfg.image.width, cfg.image.height))
# -----------------------------------------
# choose start frame point
if curTime < 10:
continue
# choose fps
if curTime % 1 != 0:
continue
# -----------------------------------------
if frameL is None:
break
# -----------------------------------------
# -----------------------------------------
# get bbox of objects
detections, scores, classes, num = odapi.processFrame(frameL)
# get bbox of persons
mask = np.array([
a & b for a, b in zip(
np.array(classes) == 1,
np.array(scores) > cfg.detector.threshold)
])
detections = np.array(detections)
if np.sum(mask) != 0:
detections = detections[mask]
else:
detections = []
# apply ion filter
detections_iou = []
for slow_loop_idx in range(len(detections)):
if detections[slow_loop_idx][0] == -1:
continue
for fast_loop_idx in range(len(detections)):
if detections[fast_loop_idx][0] == -1:
continue
r = bb_intersection_over_union(detections[slow_loop_idx],
detections[fast_loop_idx])
if r > cfg.detector.iou_threshold and slow_loop_idx != fast_loop_idx:
detections[fast_loop_idx][0] = -1
for box in detections:
if box[0] == -1:
continue
detections_iou.append(box)
# save detections after filter
detections = detections_iou
# -----------------------------------------
# start tracking part
if len(detections) != 0:
# init body_bank
result_arr = [] # ids of detected persons
result_features = [
] # features of detected persons respectively to result_arr idxes
result_dist_arr = [
] # distance of detected persons respectively to result_arr idxes
body_bank_tmp = body_bank.copy(
) # body bank copy for next loop to exclude person after search
len_body_bank_tmp = len(
body_bank_tmp) # decreasing length of body_bank_tmp
# search loop
for i, body_detection_xy in enumerate(detections):
# if body bank (gallary) is not empty
if len(body_bank) != 0:
# if tmp body bank (gallary) is empty but we still have detected persons
if len_body_bank_tmp == 0:
# get coordinates of person
p = detections[i]
p = check_coords(p)
# do preparations for image to fit it into reid network
img = np.transpose(
frameL[int(p[0]):int(p[2]),
int(p[1]):int(p[3])]).astype('f') / 255.
img = np.expand_dims(img, axis=0)
# extract features
features_img_query = F.normalize(
extract_features(torch.from_numpy(
img).cuda())).cpu().numpy()[0]
# add features of new pearson to gallary
body_bank.append(features_img_query)
result = len(body_bank)
result_arr.append(result)
#body_bank_bb.append(detections[i])
# init distance of new person
body_bank_dist.append(0)
continue
# get distance of person to persons in gallery
distmat_arr, query_f = get_dist(
[np.array(body_detection_xy)], frameL,
body_bank_tmp)
# get id of person in gallery
result = np.argmin(distmat_arr, 1)[0]
#print(distmat_arr[0][result])
# if minimum distance is lower cfg.model.threshold it is the same person
if (distmat_arr[0][result] < cfg.model.threshold):
# add id to result array
result_arr.append(result)
# exclude person from search
body_bank_tmp[result] = []
# decrease lenght of tmp body bank
len_body_bank_tmp -= 1
# save features and distance
result_features.append(query_f[result])
result_dist_arr.append(np.min(distmat_arr, 1)[0])
else:
# save features and distance
body_bank.append(query_f[result])
body_bank_dist.append(0)
else:
print("add new user")
#result = len(body_bank) + 1
# get coordinates of person
p = detections[i]
p = check_coords(p)
# do preparations for image to fit it into reid network
img = np.transpose(
frameL[int(p[0]):int(p[2]),
int(p[1]):int(p[3])]).astype('f') / 255.
img = np.expand_dims(img, axis=0)
# extract features
features_img_query = F.normalize(
extract_features(
torch.from_numpy(img).cuda())).cpu().numpy()[0]
# save features and distance
body_bank.append(features_img_query)
result = len(body_bank)
result_arr.append(result)
#body_bank_bb.append(detections[i])
body_bank_dist.append(0)
# ax for visualisation of features
ax = [0] * len(result_arr)
if cfg.visualisation.key:
if len(body_bank) != 0:
fig, ax = plt.subplots(
nrows=len(body_bank),
ncols=3,
gridspec_kw={'width_ratios': [2, 10, 10]},
figsize=(10, 2))
if len(body_bank) == 1:
ax = [ax]
for i, pack in enumerate(zip(result_arr, result_dist_arr, ax)):
id, result, row = pack
# add new human if % is too low
# refresh bank
if cfg.visualisation.key:
if len(body_bank) != 0:
for idx, col in enumerate(row):
if idx == 0:
p = detections[i]
p = check_coords(p)
col.imshow(frameL[int(p[0]):int(p[2]),
int(p[1]):int(p[3])])
if idx == 1:
col.plot(range(0,
len(body_bank[id]) + 1),
list(body_bank[id]) + [0.40])
if idx == 2:
print(
int(
"".join(
str(int(x)) for x in [
x > 0.1 and y > 0.1
for x, y in zip(
body_bank[id],
result_features[i])
]), 2))
# refresh image in gallery if image is good
if result < cfg.model.threshold:
p = detections[i]
p = check_coords(p)
body_bank_dist[
id] = result # + body_bank_dist[id]) / 2
body_bank[id] = result_features[
i] #+ body_bank[id]) / 2 #frameL[int(p[1]):int(p[3]), int(p[0]):int(p[2])]
text = "_ID_{} <{}>".format(str(id), str(round(result, 3)))
cv2.putText(frameL, text, (int(detections[i][1]) - 10,
int(detections[i][0]) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
if cfg.visualisation.key:
plt.show()
plt.pause(4)
plt.close()
frameL = frameL / 255
for p in detections:
cv2.rectangle(frameL, (p[1], p[0]), (p[3], p[2]),
(50, 50, 250), 2)
cv2.imshow("L", frameL)
#print("shape ", frameL.shape)
if args_save:
video_writer.write((frameL * 255).astype(np.uint8))
cv2.waitKey(1)
