def main():
# load config
cfg.merge_from_file(args.config)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(
torch.load(args.snapshot,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# build tracker
tracker = build_tracker(model)
first_frame = True
if args.video_name:
video_name = args.video_name.split('/')[-1].split('.')[0]
else:
video_name = 'webcam'
# cv2.namedWindow(video_name, cv2.WND_PROP_FULLSCREEN)
i = 0
for frame in get_frames(args.video_name):
if first_frame:
try:
init_rect = cv2.selectROI(video_name, frame, False, False)
except:
exit()
tracker.init(frame, init_rect)
first_frame = False
else:
outputs = tracker.track(frame)
if 'polygon' in outputs:
polygon = np.array(outputs['polygon']).astype(np.int32)
cv2.polylines(frame, [polygon.reshape((-1, 1, 2))], True,
(0, 255, 0), 3)
mask = ((outputs['mask'] > cfg.TRACK.MASK_THERSHOLD) * 255)
mask = mask.astype(np.uint8)
mask = np.stack([mask, mask * 255, mask]).transpose(1, 2, 0)
frame = cv2.addWeighted(frame, 0.77, mask, 0.23, -1)
else:
bbox = list(map(int, outputs['bbox']))
cv2.rectangle(frame, (bbox[0], bbox[1]),
(bbox[0] + bbox[2], bbox[1] + bbox[3]),
(0, 255, 0), 3)
# cv2.imshow(video_name, frame)
# cv2.waitKey(40)
print(i)
cv2.imwrite(filename="/home/tempuser1/pysot/demo/ouput/" + str(i) +
'.jpg',
img=frame)
i += 1
def main():
rank, world_size = dist_init()
logger.info("init done")
# load cfg
cfg.merge_from_file(args.cfg)
if rank == 0:
if not os.path.exists(cfg.TRAIN.LOG_DIR):
os.makedirs(cfg.TRAIN.LOG_DIR)
init_log('global', logging.INFO)
if cfg.TRAIN.LOG_DIR:
add_file_handler('global',
os.path.join(cfg.TRAIN.LOG_DIR, 'logs.txt'),
logging.INFO)
logger.info("Version Information: \n{}\n".format(commit()))
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
# create model
model = ModelBuilder().cuda().train()
dist_model = DistModule(model)
# load pretrained backbone weights
if cfg.BACKBONE.PRETRAINED:
cur_path = os.path.dirname(os.path.realpath(__file__))
backbone_path = os.path.join(cur_path, '../', cfg.BACKBONE.PRETRAINED)
load_pretrain(model.backbone, backbone_path)
# create tensorboard writer
if rank == 0 and cfg.TRAIN.LOG_DIR:
tb_writer = SummaryWriter(cfg.TRAIN.LOG_DIR)
else:
tb_writer = None
# build dataset loader
train_loader = build_data_loader()
# build optimizer and lr_scheduler
optimizer, lr_scheduler = build_opt_lr(dist_model.module,
cfg.TRAIN.START_EPOCH)
# resume training
if cfg.TRAIN.RESUME:
logger.info("resume from {}".format(cfg.TRAIN.RESUME))
assert os.path.isfile(cfg.TRAIN.RESUME), \
'{} is not a valid file.'.format(cfg.TRAIN.RESUME)
model, optimizer, cfg.TRAIN.START_EPOCH = \
restore_from(model, optimizer, cfg.TRAIN.RESUME)
dist_model = DistModule(model)
logger.info(lr_scheduler)
logger.info("model prepare done")
# start training
train(train_loader, dist_model, optimizer, lr_scheduler, tb_writer)
def main():
# load config
cfg.merge_from_file(args.config)
# create model
model = ModelBuilder()
# load model
checkpoint = torch.load(args.snapshot, map_location=torch.device('cpu'))
model.load_state_dict(checkpoint)
for param in model.parameters():
param.requires_grad = False
def build_model(self):
model = ModelBuilder()
# load model
model.load_state_dict(torch.load(args.snapshot,
map_location=lambda storage, loc: storage.cpu()))
# import ipdb
# ipdb.set_trace()
device = torch.device('cuda:{}'.format(int(self.index//self.num_workers)) if cfg.CUDA else 'cpu')
print(device)
model.eval().to(device)
# build tracker
tracker = build_tracker(model)
return tracker
def save_siamese_rpn():
# load config
cfg.merge_from_file(args.config)
cfg.BACKBONE.TYPE = 'alexnetlegacy'
# create model
model_legacy = ModelBuilder()
# load model
model_legacy = load_pretrain(model_legacy, args.snapshot).cuda().eval()
cfg.BACKBONE.TYPE = 'alexnet'
# create model
model_alexnet = ModelBuilder()
#
# for key ,item in model.named_parameters():
# print(key,item.shape)
for key, item in model_alexnet.named_parameters():
print(key, item.shape)
name_map = {}
model_legacy_dict = model_legacy.state_dict()
model_alexnet_dict = model_alexnet.state_dict()
for para1, para2 in zip(model_legacy.named_parameters(),
model_alexnet.named_parameters()):
# print(para1[0],para1[1].shape)
print(para1[0])
print(para2[0])
print(para1[1].shape)
print(para2[1].shape)
print("--" * 40)
# print("['{}'--->'{}']".format(para1[0], para2[0]),para1[1].shape, para2[1].shape)
name_map[para1[0]] = para2[0]
print(name_map)
for key, val in name_map.items():
model_alexnet_dict[val] = model_legacy_dict[key]
torch.save(model_alexnet_dict, "siamese_alexnet_rpn.pth")
def load_tracker(self, tracker_config, tracker_snapshot):
"""Load the selected pysot tracker.
Args:
- tracker_config (str): Path to pysot config file for the tracker
- tracker_snapshot (str): Path to .pth file of pysot tracker
"""
cfg.merge_from_file(tracker_config)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
model = ModelBuilder()
model.load_state_dict(torch.load(tracker_snapshot))
model.eval().to(device)
self.tracker = build_tracker(model)
def __init__(self):
super(DROL, self).__init__("DROL")
# load config
cfg.merge_from_file(path_config.DROL_CONFIG)
seed_torch(cfg.TRACK.SEED)
# create model
model = ModelBuilder()
# load model
model = load_pretrain(model, path_config.DROL_SNAPSHOT).cuda().eval()
# build tracker
self.tracker = build_tracker(model)
def main():
# load config
cfg.merge_from_file(args.config)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(torch.load(args.snapshot,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# build tracker
tracker = build_tracker(model)
first_frame = True
if args.video_name:
video_name = args.video_name.split('/')[-1].split('.')[0]
else:
video_name = 'webcam'
cv2.namedWindow(video_name, cv2.WINDOW_NORMAL)#cv2.WND_PROP_FULLSCREEN)
for frame in get_frames(args.video_name):
if first_frame:
try:
init_rect = cv2.selectROI(video_name, frame, False, False)#choose a rectangle as ROI
except:
exit()
tracker.init(frame, init_rect)#initiating the tracker
first_frame = False # choose the ROI on the first frame and then track it on the following frames
else:
outputs = tracker.track(frame)#outputs:bbox/polygon+best_score
if 'polygon' in outputs:
polygon = np.array(outputs['polygon']).astype(np.int32)
cv2.polylines(frame, [polygon.reshape((-1, 1, 2))],#draw polygons([vertex_nums,1,2]) on the frame
True, (0, 255, 0), 3)
mask = ((outputs['mask'] > cfg.TRACK.MASK_THERSHOLD) * 255)
mask = mask.astype(np.uint8)
mask = np.stack([mask, mask*255, mask]).transpose(1, 2, 0)
frame = cv2.addWeighted(frame, 0.77, mask, 0.23, -1)#image fusion, can adjust transparency
else:
bbox = list(map(int, outputs['bbox']))#float to int
cv2.rectangle(frame, (bbox[0], bbox[1]),#draw bbox on the frame
(bbox[0]+bbox[2], bbox[1]+bbox[3]),
(0, 255, 0), 3)
cv2.imshow(video_name, frame)
cv2.waitKey(40)
def __init__(self, lr_u=0.2,lr_v=0.2,lambda_u=0.1,lambda_v=10.0,x_padding=0.5, z_ratio=1.2,features='gray', kernel='gaussian'):
super(SFKCF).__init__()
self.x_padding = x_padding
self.lambda_ = 1e-4
self.features = features
self.w2c=None
if self.features=='hog':
self.interp_factor = 0.02
self.sigma = 0.5
self.cell_size=4
self.output_sigma_factor=0.1
elif self.features=='sfres50':
self.interp_factor = 0.02
self.sigma = 0.5
self.cell_size=8.0
self.output_sigma_factor=0.1
model = ModelBuilder()
model = load_pretrain(model, cfg.BACKBONE.PRETRAINED).backbone
self.model = model.cuda().eval()
elif self.features=='gray' or self.features=='color':
self.interp_factor=0.075
self.sigma=0.2
self.cell_size=1
self.output_sigma_factor=0.1
elif self.features=='cn':
self.interp_factor=0.075
self.sigma=0.2
self.cell_size=1
self.output_sigma_factor=1./16
self.padding=1
else:
raise NotImplementedError
self.kernel=kernel
self.U = None
self.V = None
self.lr_u = lr_u
self.lr_v = lr_v
self.lambda_v = lambda_v
self.lambda_u = lambda_u
self.z_padding = z_ratio*x_padding
self.vis = None
def load_pysot_model(tracker_type):
configpath = "./week3/kalman/pysot/experiments/" + PYSOT_TRACKERS[tracker_type] + \
"/config.yaml"
modelpath = "./week3/kalman/pysot/models/" + PYSOT_TRACKERS[
tracker_type] + ".pth"
cfg.merge_from_file(configpath)
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# load model
model = ModelBuilder()
model.load_state_dict(
torch.load(modelpath, map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
return load_pretrain(model, modelpath).cuda().eval()
def __init__(self,dataset=''):
if 'OTB' in dataset:
cfg_file = os.path.join(project_path_,'pysot/experiments/siamrpn_r50_l234_dwxcorr_otb/config.yaml')
snapshot = os.path.join(project_path_,'pysot/experiments/siamrpn_r50_l234_dwxcorr_otb/model.pth')
elif 'LT' in dataset:
cfg_file = os.path.join(project_path_, 'pysot/experiments/siamrpn_r50_l234_dwxcorr_lt/config.yaml')
snapshot = os.path.join(project_path_, 'pysot/experiments/siamrpn_r50_l234_dwxcorr_lt/model.pth')
else:
cfg_file = os.path.join(project_path_, 'pysot/experiments/siamrpn_r50_l234_dwxcorr/config.yaml')
snapshot = os.path.join(project_path_, 'pysot/experiments/siamrpn_r50_l234_dwxcorr/model.pth')
# load config
cfg.merge_from_file(cfg_file)
# create model
self.model = ModelBuilder()# A Neural Network.(a torch.nn.Module)
# load model
self.model = load_pretrain(self.model, snapshot).cuda().eval()
def PYSOTINIT():
# load config
cfg.merge_from_file(tracker_config)
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(
torch.load(snapshot, map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# build tracker
tracker = build_tracker(model)
return tracker
def __init__(self, config_file, model_file):
self.config_file = config_file
self.model_file = model_file
# load config
cfg.merge_from_file(self.config_file)
cfg.CUDA = torch.cuda.is_available()
self.device = torch.device('cuda' if cfg.CUDA else 'cpu')
# load model
self.model = ModelBuilder()
self.model.load_state_dict(
torch.load(model_file,
map_location=lambda storage, loc: storage.cpu()))
self.model.eval().to(self.device)
# build tracker
self.tracker = build_tracker(self.model)
def __init__(self, config, snapshot):
cfg.merge_from_file(config)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
self.model = ModelBuilder()
# load model
self.model.load_state_dict(
torch.load(snapshot,
map_location=lambda storage, loc: storage.cpu()))
self.model.eval().to(device)
# build tracker
self.tracker = build_tracker(self.model)
self.center_pos = None
self.size = None
def __init__(self, backbone, target):
super(SiamRPNPPGroup, self).__init__(f"SiamRPN++Group/{backbone}/{target}")
if backbone == "AlexNet" and target == "OTB":
config = path_config.SIAMRPNPP_ALEXNET_OTB_CONFIG
snapshot = path_config.SIAMRPNPP_ALEXNET_OTB_SNAPSHOT
elif backbone == "AlexNet" and target == "VOT":
config = path_config.SIAMRPNPP_ALEXNET_CONFIG
snapshot = path_config.SIAMRPNPP_ALEXNET_SNAPSHOT
elif backbone == "ResNet-50" and target == "OTB":
config = path_config.SIAMRPNPP_RESNET_OTB_CONFIG
snapshot = path_config.SIAMRPNPP_RESNET_OTB_SNAPSHOT
elif backbone == "ResNet-50" and target == "VOT":
config = path_config.SIAMRPNPP_RESNET_CONFIG
snapshot = path_config.SIAMRPNPP_RESNET_SNAPSHOT
elif backbone == "ResNet-50" and target == "VOTLT":
config = path_config.SIAMRPNPP_RESNET_LT_CONFIG
snapshot = path_config.SIAMRPNPP_RESNET_LT_SNAPSHOT
elif backbone == "MobileNetV2" and target == "VOT":
config = path_config.SIAMRPNPP_MOBILENET_CONFIG
snapshot = path_config.SIAMRPNPP_MOBILENET_SNAPSHOT
elif backbone == "SiamMask" and target == "VOT":
config = path_config.SIAMPRNPP_SIAMMASK_CONFIG
snapshot = path_config.SIAMPRNPP_SIAMMASK_SNAPSHOT
else:
raise ValueError("Invalid backbone and target")
# load config
cfg.merge_from_file(config)
cfg.CUDA = torch.cuda.is_available()
device = torch.device("cuda" if cfg.CUDA else "cpu")
# create model
self.model = ModelBuilder()
# load model
self.model.load_state_dict(
torch.load(snapshot, map_location=lambda storage, loc: storage.cpu())
)
self.model.eval().to(device)
# build tracker
self.tracker = build_tracker(self.model)
def init_track(self):
# 参数整合
cfg.merge_from_file(self.config_path)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
self.textBws_show_process.append('模型对象创建...')
self.checkpoint=torch.load(self.snapshot_path, map_location=lambda storage, loc: storage.cpu())
self.model = ModelBuilder()
print('断点')
# load model
self.model.load_state_dict(self.checkpoint)
self.model.eval().to(device)
self.textBws_show_process.append('加载跟踪模型完毕!')
# 创建跟踪器
self.tracker = build_tracker(self.model)
def test_snapshot(epoch: int, snapshot: str, test_path: str):
# model
max_img = 8
model = ModelBuilder()
data = torch.load(snapshot,
map_location=lambda storage, loc: storage.cpu())
model.load_state_dict(data['state_dict'])
model.eval().to(torch.device('cpu'))
tracker = build_tracker(model)
root = cfg.DATASET.COCO.ROOT
cur_path = os.path.dirname(os.path.realpath(__file__))
root = os.path.join(cur_path, '../../', root)
anno_path = os.path.join(root, '../', "val2017.json")
with open(anno_path, 'r') as f:
anno = json.load(f)
anno = filter_zero(anno)
dataset = os.path.join(root, "val2017")
folder = random.choice(glob.glob(f"{dataset}/**"))
zs = glob.glob(f"{folder}/*.z.jpg")
xs = glob.glob(f"{folder}/*.x.jpg")
zs = sorted(zs)
xs = sorted(xs)
xs = [(x, get_anno_from_img_path(anno, x)) for x in xs]
for i in range(len(zs[:max_img])):
z = cv2.imread(zs[i])
x_path, bbox = xs[i]
x = cv2.imread(x_path)
tracker.init_(z)
cls, (x1, y1, x2, y2) = tracker.track(x)
cv2.rectangle(x, (x1, y1), (x2, y2), (255, 0, 0), 2)
a1, b1, a2, b2 = bbox
cv2.rectangle(x, (a1, b1), (a2, b2), (0, 0, 255), 2)
cv2.putText(x, 'Acc: ' + cls.astype('str'), (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
parent_dir = f"{test_path}/{os.path.basename(Path(zs[i]).parent)}"
if not os.path.exists(parent_dir):
os.makedirs(parent_dir)
cv2.imwrite(f"{parent_dir}/{os.path.basename(x_path)}", x)
cv2.imwrite(f"{parent_dir}/{os.path.basename(zs[i])}", z)
def __init__(self):
self.init_rect = None
self.pysot_pub = rospy.Publisher(config.TRACK_PUB_TOPIC,
Int32MultiArray,
queue_size=10)
self.img_sub = rospy.Subscriber(config.IMAGE_SUB_TOPIC, Image,
self.receive_frame_and_track)
self.service = rospy.Service("init_rect", InitRect, self.set_init_rect)
cfg.TRACK.TYPE = config.TRACK_TYPE
cfg.merge_from_file(config.CONFIG_PATH)
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
model = ModelBuilder()
model.load_state_dict(
torch.load(config.MODEL_PATH,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
self.tracker = build_tracker(model)
def __init__(self):
super(SiamRPNPP, self).__init__("SiamRPN++")
config = path_config.SIAMRPNPP_CONFIG
snapshot = path_config.SIAMRPNPP_SNAPSHOT
# load config
cfg.merge_from_file(config)
cfg.CUDA = torch.cuda.is_available()
device = torch.device("cuda" if cfg.CUDA else "cpu")
# create model
self.model = ModelBuilder()
# load model
self.model.load_state_dict(
torch.load(snapshot, map_location=lambda storage, loc: storage.cpu())
)
self.model.eval().to(device)
# build tracker
self.tracker = build_tracker(self.model)
def main():
# load config
cfg.merge_from_file(args.config)
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model = load_pretrain(model, args.snapshot).eval().to(device)
# build tracker
tracker = SiamAPNTracker(model, cfg.TRACK)
hp = {'lr': 0.3, 'penalty_k': 0.04, 'window_lr': 0.4}
first_frame = True
if args.video_name:
video_name = args.video_name.split('/')[-1].split('.')[0]
else:
video_name = 'webcam'
cv2.namedWindow(video_name, cv2.WND_PROP_FULLSCREEN)
for frame in get_frames(args.video_name):
if first_frame:
try:
init_rect = cv2.selectROI(video_name, frame, False, False)
except:
exit()
tracker.init(frame, init_rect)
first_frame = False
else:
outputs = tracker.track(frame, hp)
bbox = list(map(int, outputs['bbox']))
cv2.rectangle(frame, (bbox[0], bbox[1]),
(bbox[0] + bbox[2], bbox[1] + bbox[3]), (0, 255, 0),
3)
cv2.imshow(video_name, frame)
cv2.waitKey(40)
def init_track(self):
# 配置config文件
config_path = './models/siamrpn_alex_dwxcorr/config.yaml'
# 配置snapshot 文件
snapshot_path = './models/siamrpn_alex_dwxcorr/model.pth'
# 参数整合
cfg.merge_from_file(config_path)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(
torch.load(snapshot_path,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# 创建跟踪器
self.tracker = build_tracker(model)
def __init__(self, parent=None):
super(MyMainWindow, self).__init__(parent)
self.isDracula = False
# Connect the on-clicked functions
self.pushButton_locationLoading.clicked.connect(self.location_loading)
self.pushButton_videoLoading.clicked.connect(self.video_loading)
self.pushButton_cameraLoading.clicked.connect(self.camera_loading)
self.pushButton_bboxSetting.clicked.connect(self.bbox_setting)
self.pushButton_algorithmProcessing.clicked.connect(
self.algorithm_processing)
self.scrollBar.valueChanged.connect(self.slider_change)
self.selectBox.valueChanged.connect(self.select_change)
self.checkBox.stateChanged.connect(self.checkbox_change)
# Message box ignore
self.bbox_tips = True
self.save_tips = True
# Initialize trackers
model_location = './pysot/experiments/siammaske_r50_l3'
self.config = model_location + '/config.yaml'
self.snapshot = model_location + '/model.pth'
self.tracker_name = model_location.split('/')[-1]
self.video_name = ''
cfg.merge_from_file(self.config)
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
model = ModelBuilder()
model.load_state_dict(
torch.load(self.snapshot,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
self.tracker = build_tracker(model)
self.vs = None
self.analysis_box = None
self.analysis_max = 10
self.save_location = ''
self.afterCamera = False
self.bbox_list_predict = [] # [time][tracker]
def main():
# load config
cfg.merge_from_file(args.config)
cur_dir = os.path.dirname(os.path.realpath(__file__))
dataset_root = os.path.join(cur_dir, '../testing_dataset', args.dataset)
# create model
model = ModelBuilder(cfg)
# load model
model = load_pretrain(model, args.snapshot).cuda().eval()
# build tracker
tracker = build_tracker(model)
# create dataset
dataset = DatasetFactory.create_dataset(name=args.dataset,
dataset_root=dataset_root,
load_img=False)
model_name = args.snapshot.split('/')[-1].split('.')[0]
total_lost = 0
if args.dataset in ['VOT2016', 'VOT2018', 'VOT2019']:
# restart tracking
for v_idx, video in enumerate(dataset):
if args.video != '':
# test one special video
if video.name != args.video:
continue
frame_counter = 0
lost_number = 0
toc = 0
pred_bboxes = []
for idx, (img, gt_bbox) in enumerate(video):
if len(gt_bbox) == 4:
gt_bbox = [gt_bbox[0], gt_bbox[1],
gt_bbox[0], gt_bbox[1]+gt_bbox[3]-1,
gt_bbox[0]+gt_bbox[2]-1, gt_bbox[1]+gt_bbox[3]-1,
gt_bbox[0]+gt_bbox[2]-1, gt_bbox[1]]
tic = cv2.getTickCount()
if idx == frame_counter:
cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_bbox))
gt_bbox_ = [cx-(w-1)/2, cy-(h-1)/2, w, h]
tracker.init(img, gt_bbox_)
pred_bbox = gt_bbox_
pred_bboxes.append(1)
elif idx > frame_counter:
outputs = tracker.track(img)
pred_bbox = outputs['bbox']
if cfg.MASK.MASK:
pred_bbox = outputs['polygon']
overlap = vot_overlap(pred_bbox, gt_bbox, (img.shape[1], img.shape[0]))
if overlap > 0:
# not lost
pred_bboxes.append(pred_bbox)
else:
# lost object
pred_bboxes.append(2)
frame_counter = idx + 5 # skip 5 frames
lost_number += 1
else:
pred_bboxes.append(0)
toc += cv2.getTickCount() - tic
if idx == 0:
cv2.destroyAllWindows()
if args.vis and idx > frame_counter:
cv2.polylines(img, [np.array(gt_bbox, np.int).reshape((-1, 1, 2))],
True, (0, 255, 0), 3)
if cfg.MASK.MASK:
cv2.polylines(img, [np.array(pred_bbox, np.int).reshape((-1, 1, 2))],
True, (0, 255, 255), 3)
else:
bbox = list(map(int, pred_bbox))
cv2.rectangle(img, (bbox[0], bbox[1]),
(bbox[0]+bbox[2], bbox[1]+bbox[3]), (0, 255, 255), 3)
cv2.putText(img, str(idx), (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
cv2.putText(img, str(lost_number), (40, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
cv2.imshow(video.name, img)
cv2.waitKey(1)
toc /= cv2.getTickFrequency()
# save results
video_path = os.path.join('results', args.dataset, model_name,
'baseline', video.name)
if not os.path.isdir(video_path):
os.makedirs(video_path)
result_path = os.path.join(video_path, '{}_001.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in pred_bboxes:
if isinstance(x, int):
f.write("{:d}\n".format(x))
else:
f.write(','.join([vot_float2str("%.4f", i) for i in x])+'\n')
print('({:3d}) Video: {:12s} Time: {:4.1f}s Speed: {:3.1f}fps Lost: {:d}'.format(
v_idx+1, video.name, toc, idx / toc, lost_number))
total_lost += lost_number
print("{:s} total lost: {:d}".format(model_name, total_lost))
else:
# OPE tracking
for v_idx, video in enumerate(dataset):
if args.video != '':
# test one special video
if video.name != args.video:
continue
toc = 0
pred_bboxes = []
scores = []
track_times = []
for idx, (img, gt_bbox) in enumerate(video):
tic = cv2.getTickCount()
if idx == 0:
cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_bbox))
gt_bbox_ = [cx-(w-1)/2, cy-(h-1)/2, w, h]
tracker.init(img, gt_bbox_)
pred_bbox = gt_bbox_
scores.append(None)
if 'VOT2018-LT' == args.dataset:
pred_bboxes.append([1])
else:
pred_bboxes.append(pred_bbox)
else:
outputs = tracker.track(img)
pred_bbox = outputs['bbox']
pred_bboxes.append(pred_bbox)
scores.append(outputs['best_score'])
toc += cv2.getTickCount() - tic
track_times.append((cv2.getTickCount() - tic)/cv2.getTickFrequency())
if idx == 0:
cv2.destroyAllWindows()
if args.vis and idx > 0:
gt_bbox = list(map(int, gt_bbox))
pred_bbox = list(map(int, pred_bbox))
cv2.rectangle(img, (gt_bbox[0], gt_bbox[1]),
(gt_bbox[0]+gt_bbox[2], gt_bbox[1]+gt_bbox[3]), (0, 255, 0), 3)
cv2.rectangle(img, (pred_bbox[0], pred_bbox[1]),
(pred_bbox[0]+pred_bbox[2], pred_bbox[1]+pred_bbox[3]), (0, 255, 255), 3)
cv2.putText(img, str(idx), (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
cv2.imshow(video.name, img)
cv2.waitKey(1)
toc /= cv2.getTickFrequency()
# save results
if 'VOT2018-LT' == args.dataset:
video_path = os.path.join('results', args.dataset, model_name,
'longterm', video.name)
if not os.path.isdir(video_path):
os.makedirs(video_path)
result_path = os.path.join(video_path,
'{}_001.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in pred_bboxes:
f.write(','.join([str(i) for i in x])+'\n')
result_path = os.path.join(video_path,
'{}_001_confidence.value'.format(video.name))
with open(result_path, 'w') as f:
for x in scores:
f.write('\n') if x is None else f.write("{:.6f}\n".format(x))
result_path = os.path.join(video_path,
'{}_time.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in track_times:
f.write("{:.6f}\n".format(x))
elif 'GOT-10k' == args.dataset:
video_path = os.path.join('results', args.dataset, model_name, video.name)
if not os.path.isdir(video_path):
os.makedirs(video_path)
result_path = os.path.join(video_path, '{}_001.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in pred_bboxes:
f.write(','.join([str(i) for i in x])+'\n')
result_path = os.path.join(video_path,
'{}_time.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in track_times:
f.write("{:.6f}\n".format(x))
else:
model_path = os.path.join('results', args.dataset, model_name)
if not os.path.isdir(model_path):
os.makedirs(model_path)
result_path = os.path.join(model_path, '{}.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in pred_bboxes:
f.write(','.join([str(i) for i in x])+'\n')
print('({:3d}) Video: {:12s} Time: {:5.1f}s Speed: {:3.1f}fps'.format(
v_idx+1, video.name, toc, idx / toc))
else:
raise NotImplementedError
if __name__ == '__main__':
torch.backends.cudnn.benchmark = True
cfg.merge_from_file(args.config)
cur_dir = os.path.dirname(os.path.realpath(__file__))
dataset_root = os.path.join(cur_dir, '../testing_dataset', args.dataset)
dataset = DatasetFactory.create_dataset(name=args.dataset,
dataset_root=dataset_root,
load_img=False)
benchmark = EAOBenchmark(dataset)
model = ModelBuilder()
model = load_pretrain(model, args.snapshot).cuda().eval()
# the resources you computer have, object_store_memory is shm
ray.init(num_gpus=1, num_cpus=8, object_store_memory=30000000000)
tune.register_trainable('fitness', fitness)
# define search space
params = {
'penalty_k': hp.quniform('penalty_k', 0.001, 0.6, 0.001),
'lr': hp.quniform('scale_lr', 0.1, 0.8, 0.001),
'window_influence': hp.quniform('window_influence', 0.05, 0.65, 0.001),
'search_region': hp.choice('search_region', [255]),
}
# stop condition for VOT and OTB
def main():
is_gpu_cuda_available = torch.cuda.is_available()
if not is_gpu_cuda_available:
raise RuntimeError(
'Failed to locate a CUDA GPU. Program cannot continue..')
num_gpus = torch.cuda.device_count()
gpu_type = torch.cuda.get_device_name(0)
print(f"You have {num_gpus} available of type: {gpu_type}")
print("This might take a few minutes...Grab a cup of coffee\n")
# load config
cfg.merge_from_file(args.config)
dataset_root = os.path.join(args.dataset_directory, args.dataset)
print(f"dataset root-->{dataset_root}")
# create model
model = ModelBuilder()
# load model
model = load_pretrain(model, args.snapshot).cuda().eval()
# build tracker
tracker = build_tracker(model)
# create dataset
dataset = DatasetFactory.create_dataset(name=args.dataset,
dataset_root=dataset_root,
load_img=False)
model_name = args.model_name
print(f"Model name is {model_name}")
total_lost = 0
if args.dataset in vot_like_dataset:
# restart tracking
for v_idx, video in enumerate(dataset):
if args.video != '':
# test one special video
if video.name != args.video:
continue
frame_counter = 0
lost_number = 0
toc = 0
pred_bboxes = []
for idx, (img, gt_bbox) in enumerate(video):
if len(gt_bbox) == 4:
gt_bbox = [
gt_bbox[0], gt_bbox[1], gt_bbox[0],
gt_bbox[1] + gt_bbox[3] - 1,
gt_bbox[0] + gt_bbox[2] - 1,
gt_bbox[1] + gt_bbox[3] - 1,
gt_bbox[0] + gt_bbox[2] - 1, gt_bbox[1]
]
tic = cv2.getTickCount()
if idx == frame_counter:
cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_bbox))
gt_bbox_ = [cx - (w - 1) / 2, cy - (h - 1) / 2, w, h]
tracker.init(img, gt_bbox_)
pred_bbox = gt_bbox_
pred_bboxes.append(1)
elif idx > frame_counter:
outputs = tracker.track(img)
pred_bbox = outputs['bbox']
if cfg.MASK.MASK:
pred_bbox = outputs['polygon']
overlap = vot_overlap(pred_bbox, gt_bbox,
(img.shape[1], img.shape[0]))
if overlap > 0.85:
# not lost
pred_bboxes.append(pred_bbox)
else:
# lost object
pred_bboxes.append(2)
frame_counter = idx + args.skip_frames # skip 1 frame
lost_number += 1
else:
pred_bboxes.append(0)
toc += cv2.getTickCount() - tic
if idx == 0:
cv2.destroyAllWindows()
if args.vis and idx > frame_counter:
cv2.polylines(
img, [np.array(gt_bbox, np.int).reshape(
(-1, 1, 2))], True, (0, 255, 0), 3)
if cfg.MASK.MASK:
cv2.polylines(
img,
[np.array(pred_bbox, np.int).reshape(
(-1, 1, 2))], True, (0, 255, 255), 3)
else:
bbox = list(map(int, pred_bbox))
cv2.rectangle(img, (bbox[0], bbox[1]),
(bbox[0] + bbox[2], bbox[1] + bbox[3]),
(0, 255, 255), 3)
cv2.putText(img, str(idx), (40, 40),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
cv2.putText(img, str(lost_number), (40, 80),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
cv2.imshow(video.name, img)
cv2.waitKey(1)
toc /= cv2.getTickFrequency()
# save results
save_path = os.path.join(args.results_path, args.dataset,
model_name, args.experiment_name,
video.name)
if not os.path.isdir(save_path):
os.makedirs(save_path)
result_path = os.path.join(save_path,
'{}_001.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in pred_bboxes:
if isinstance(x, int):
f.write("{:d}\n".format(x))
else:
f.write(','.join([vot_float2str("%.4f", i)
for i in x]) + '\n')
with open(os.path.join(save_path, '..', 'lost.txt'), 'a+') as f:
f.write(
f"{v_idx+1} Class: {video.name} | Time: {toc}s | Speed: {idx/toc}fps | Lost:{lost_number} \n"
)
print(
'({:3d}) Class: {:12s} Time: {:4.1f}s Speed: {:3.1f}fps Lost: {:d}'
.format(v_idx + 1, video.name, toc, idx / toc, lost_number))
total_lost += lost_number
print("{:s} total lost: {:d}".format(model_name, total_lost))
with open(os.path.join(save_path, '..', 'lost.txt'), 'a+') as f:
f.write(
f"Model architeture used --> {model_name} \ntotal lost: {total_lost} \n"
)
f.write(f"SKIP FRAMES USED --> {args.skip_frames}")
else:
# OPE tracking
# will be implemented if needed in future
pass
return cv2.LUT(image, table)
color_img = np.zeros((1280, 720, 3), dtype=np.uint8)
result_mask_img = np.zeros((1280, 720, 3), dtype=np.uint8)
result_bbox_img = np.zeros((1280, 720, 3), dtype=np.uint8)
result_mask = np.zeros((1280, 720), dtype=np.uint8)
pysot_img = np.zeros((1280, 720, 3), dtype=np.uint8)
mask_rcnn_flag = 0
pysot_mask = np.zeros((1280, 720), dtype=np.uint8)
pysot_contour_img = np.zeros((1280, 720, 3), dtype=np.uint8)
cfg.merge_from_file('config.yaml')
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
model_pysot = ModelBuilder()
tracker = build_tracker(model_pysot)
model_pysot.load_state_dict(
torch.load('model.pth', map_location=lambda storage, loc: storage.cpu()))
model_pysot.eval().to(device)
def run_maskrcnn():
global color_img
global result_mask_img
global result_bbox_img
global result_mask
global mask_rcnn_flag
global inds_len
while 1:
mask_rcnn_flag = 1
def main():
torch.cuda.set_device(args.gpu_id)
model_dir = "./experiments/siamrpn_r50_l234_dwxcorr/model.pth"
model_config = "./experiments/siamrpn_r50_l234_dwxcorr/config.yaml"
if os.path.isfile(model_dir):
print("model file {} found".format(model_dir))
else:
print("model files not found, starting download".format(model_dir))
os.system(
"gdown https://drive.google.com/uc?id=1-tEtYQdT1G9kn8HsqKNDHVqjE16F8YQH")
os.system("mv model.pth ./experiments/siamrpn_r50_l234_dwxcorr")
# load config
cfg.merge_from_file(model_config)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(torch.load(model_dir,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# create an unique identifier
worker_id = uuid.uuid4()
# build tracker
tracker = build_tracker(model)
# Socket to talk to server
context = zmq.Context()
sub_socket = context.socket(zmq.SUB)
# set up frame listening socket
sub_socket.connect("tcp://{}:5556".format(args.server_ip))
sub_socket.setsockopt_string(zmq.SUBSCRIBE, "frame_")
sub_socket.setsockopt_string(zmq.SUBSCRIBE, str(worker_id))
# setup push socket
context = zmq.Context()
push_socket = context.socket(zmq.PUSH)
push_socket.connect("tcp://{}:5557".format(args.server_ip))
# event monitoring
# used to register worker once connection is established
EVENT_MAP = {}
for name in dir(zmq):
if name.startswith('EVENT_'):
value = getattr(zmq, name)
EVENT_MAP[value] = name
# monitor thread function
def event_monitor(monitor):
while monitor.poll():
evt = recv_monitor_message(monitor)
evt.update({'description': EVENT_MAP[evt['event']]})
if evt['event'] == zmq.EVENT_HANDSHAKE_SUCCEEDED:
push_socket.send_json(
{"type": "REGISTER", "id": str(worker_id)})
if evt['event'] == zmq.EVENT_MONITOR_STOPPED:
break
monitor.close()
# register monitor
monitor = sub_socket.get_monitor_socket()
t = threading.Thread(target=event_monitor, args=(monitor,))
t.start()
support = None
try:
while True:
# wait for next message
_ = sub_socket.recv()
md = sub_socket.recv_json()
if md['type'] == 'FRAME':
msg = sub_socket.recv()
buf = memoryview(msg)
frame = np.frombuffer(
buf, dtype=md['dtype']).reshape(md['shape'])
if support is None:
continue
outputs = tracker.track(frame)
bbox = list(map(int, outputs['bbox']))
# send result
push_socket.send_json(
{
"type": "TRACK",
"bbox": bbox,
"score": outputs['best_score'].tolist(),
"time": md['time'],
"id": str(worker_id)
})
print('message: {}'.format(md['time']), end='\r')
elif md['type'] == 'SUPPORT':
frame_raw = md['data']['img'] # base 64 png image
frame = np.array(
Image.open(
io.BytesIO(
base64.b64decode(frame_raw)
)
).convert('RGB'))[:, :, ::-1]
bbox = [int(float(i)) for i in md['data']['bbox'].split(",")]
tracker.init(frame, bbox)
support = (frame, bbox)
print('Support received, tracking will now start')
elif md['type'] == 'LOCATION':
# make sure tracker has been initalized
if support is not None:
center_pos = np.array(md['data'])
tracker.update(center_pos)
elif md['type'] == 'PING':
push_socket.send_json({"type": "PONG", "id": str(worker_id)})
else:
print('Invalid message type received: {}'.format(md['type']))
except KeyboardInterrupt:
print('Exiting... notifying server of disconnect')
push_socket.send_json(
{"type": "FIN", "id": str(worker_id)})
# wait for the server to respond or let the user forcefully close
print("Waiting for server response. Press CTRL+C again to forcefully close")
while True:
_ = sub_socket.recv()
md = sub_socket.recv_json()
if md['type'] == "FIN":
print('Server responded, now exiting')
exit(0)
elif md['type'] == "FRAME":
# we have to accept the incoming frame to properly accept future messages
msg = sub_socket.recv()
def main():
cfg1 = './cfg/yolov3-spp.cfg'
data_cfg = './data/coco.data'
weights = './weights/yolov3-spp.weights'
images = './videos1'
output='output'
fourcc='mp4v'
img_size=416
conf_thres=0.5
nms_thres = 0.5
save_images=True
#yolo Initialize
device = torch_utils.select_device()
torch.backends.cudnn.benchmark = False # set False for reproducible results
if os.path.exists(output):
shutil.rmtree(output) # delete output folder
os.makedirs(output) # make new output folder
#yolo Initialize model
if ONNX_EXPORT:
s = (320, 192) # (320, 192) or (416, 256) or (608, 352) onnx model image size (height, width)
model = Darknet(cfg1, s)
else:
model = Darknet(cfg1, img_size)
#yolo Load weights
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location=device)['model'])
else: # darknet format
_ = load_darknet_weights(model, weights)
#yolo Fuse Conv2d + BatchNorm2d layers
model.fuse()
#yolo Eval mode
model.to(device).eval()
if ONNX_EXPORT:
img = torch.zeros((1, 3, s[0], s[1]))
torch.onnx.export(model, img, 'weights/export.onnx', verbose=True)
return
#Set Dataloader
vid_path, vid_writer = None, None
dataloader = LoadImages(images, img_size=img_size)
# Get classes and colors
classes = load_classes(parse_data_cfg(data_cfg)['names'])
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(classes))]
for i, (path, img, im0, vid_cap) in enumerate(dataloader):
the_list = []
t = time.time()
save_path = str(Path(output) / Path(path).name)
video_name_yolo = os.path.basename(save_path)
frame_count = 0
ignore = 15
# Get detections
frame_count = frame_count + 15
img = torch.from_numpy(img).unsqueeze(0).to(device)
pred, _ = model(img)
det = non_max_suppression(pred, conf_thres, nms_thres)[0]
if det is not None and len(det) > 0:
# Rescale boxes from 416 to true image size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
for *xyxy,conf,cls_conf,cls in det:
if cls == 0:
if conf >= 0.75:
my_list_det = ('%g,' * 4) % (*xyxy,)
my_list_list = [my_list_det]
my_list_strp = my_list_list[0][:-1]
my_list_int = np.array(my_list_strp.split(",")).astype('int').tolist()
my_int = np.array(my_list_strp.split(","))
the_list.append(my_list_int)
with open('./videos1/'+video_name_yolo+'_yolo.csv', 'a', newline='') as csvfile:
writer = csv.writer(csvfile, quoting=0, delimiter = ",")#,quotechar='',escapechar='')
writer.writerow(my_int)
else:
continue
else:
continue
# if save_images: # Save image with detections
# if dataloader.mode == 'images':
# cv2.imwrite(save_path, im0)
# else:
# if vid_path != save_path: # new video
# vid_path = save_path
# if isinstance(vid_writer, cv2.VideoWriter):
# vid_writer.release() # release previous video writer
# fps = vid_cap.get(cv2.CAP_PROP_FPS)
# width = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# height = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (width, height))
# vid_writer.write(im0)
# def remove(string):
# return string.replace(" ", ",")
# the_list_app = []
# for x in the_list:
# the_list_r = remove(x)
# the_list_app.append(the_list_r)
# print(the_list_app)
# if len(master_tracking_list) == 0:
# a = [the_list[0]]
# for track_cor in master_tracking_list:
# a = [remove(cordinate) if cordinate <= 0.20 else print('ignored') for cordinate in [bbox.IoU(track_cor,cor) for cor in the_list]]
#siam load config
cfg.merge_from_file('./experiments/siamrpn_r50_l234_dwxcorr_8gpu/config.yaml')
cfg.CUDA = torch.cuda.is_available()
#device = torch.device('cuda' if cfg.CUDA else 'cpu')
#siam create model
model = ModelBuilder()
#siam load model
model.load_state_dict(torch.load('./experiments/siamrpn_r50_l234_dwxcorr_8gpu/model.pth', map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
#siam build tracker
tracker = build_tracker(model)
#video_list = []
video_list_mp4 = glob1("./videos1", "*.mp4")
#video_list.append(video_list_mp4)
video_list_avi = glob1("./videos1","*.avi")
#video_list.append(video_list_avi)
#print(video_list)
for video_name in itertools.chain(video_list_mp4, video_list_avi):
print(video_name)
#with open('./demo/vids/'+video_name+'_yolo.csv') as vid:
#df = pd.read_csv('./demo/vids/'+video_name+'_yolo.csv', delimiter = ',', dtype=int)
# read = csv.reader(vid, delimiter = ',', newline='')
# for row in read:
# cords = row[0]
#print(cords)
#cords = [list(x) for x in df.values]
#cords = [1,2,3,4]
try:
df = pd.read_csv('./videos1/'+video_name+'_yolo.csv', delimiter=',', header=None)
cords = [list(x) for x in df.values]
except:
continue
master_tracking_list = []
object_counter = 0
for cord in cords:
cord = [cord[0],cord[1],cord[2]-cord[0],cord[3]-cord[1]]
try:
iou = [bbx.IoU(track_cor,cord) for track_cor in master_tracking_list]
except:
pass
if len(iou) == 0:
master_tracking_list.append(cord)
iou = [0.0]
if max(iou) <= 0.3:
object_counter = object_counter + 1
first_frame = True
# if video_name:#args.video_name:
# video_name = video_name.split('/')[-1].split('.')[0]
# print(video_name)
# #video_name = args.video_name.split('/')[-1].split('.')[0]
# else:
# exit()
#print(cord)
frame_count = 0
#mylist = [frame_count,object_counter,cord[0],cord[1],cord[2],cord[3],video_name]
for frame in get_frames(video_name):#(args.video_name):
if first_frame:
try:
init_rect = cord
except:
exit()
tracker.init(frame, init_rect)
first_frame = False
else:
outputs = tracker.track(frame)
if 'polygon' in outputs:
exit()
else:
#crds = map(int,outputs['bbox'])
bbox = list(map(int,outputs['bbox']))
master_tracking_list.append(bbox)
#cv2.rectangle(frame,(bbox[0],bbox[1]),(bbox[0]+bbox[2],bbox[1]+bbox[3]),(0,255,0),3)
frame_count = frame_count + 1
mylist = [frame_count,object_counter,bbox[0],bbox[1],bbox[2],bbox[3],video_name]
#cv2.imshow(video_name, frame)
#cv2.waitKey(40)
with open('./videos1/vid-'+str(video_name)+'-tracking-'+str(object_counter)+'-object-'+str(cord)+'.csv', 'a', newline='') as csvfile:
writer = csv.writer(csvfile, quoting=0)#,quotechar='',escapechar='')
writer.writerow(mylist)
else:
print("Coordinate Ignored")
continue
def main():
# load config
cfg.merge_from_file(args.config)
cur_dir = os.path.dirname(os.path.realpath(__file__))
dataset_root = os.path.join(cur_dir, '../test_dataset', args.dataset)
# create model
model = ModelBuilder()
# load model
model = load_pretrain(model, args.snapshot).cuda().eval()
# build tracker
tracker = SiamAPNTracker(model)
# create dataset
dataset = DatasetFactory.create_dataset(name=args.dataset,
dataset_root=dataset_root,
load_img=False)
model_name = args.snapshot.split('/')[-1].split('.')[0] + str(cfg.TRACK.w1)
# OPE tracking
for v_idx, video in enumerate(dataset):
if args.video != '':
# test one special video
if video.name != args.video:
continue
toc = 0
pred_bboxes = []
scores = []
track_times = []
for idx, (img, gt_bbox) in enumerate(video):
tic = cv2.getTickCount()
if idx == 0:
cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_bbox))
gt_bbox_ = [cx - (w - 1) / 2, cy - (h - 1) / 2, w, h]
tracker.init(img, gt_bbox_)
pred_bbox = gt_bbox_
scores.append(None)
if 'VOT2018-LT' == args.dataset:
pred_bboxes.append([1])
else:
pred_bboxes.append(pred_bbox)
else:
outputs = tracker.track(img)
pred_bbox = outputs['bbox']
pred_bboxes.append(pred_bbox)
scores.append(outputs['best_score'])
toc += cv2.getTickCount() - tic
track_times.append(
(cv2.getTickCount() - tic) / cv2.getTickFrequency())
if idx == 0:
cv2.destroyAllWindows()
if args.vis and idx > 0:
gt_bbox = list(map(int, gt_bbox))
pred_bbox = list(map(int, pred_bbox))
cv2.rectangle(
img, (gt_bbox[0], gt_bbox[1]),
(gt_bbox[0] + gt_bbox[2], gt_bbox[1] + gt_bbox[3]),
(0, 255, 0), 3)
cv2.rectangle(
img, (pred_bbox[0], pred_bbox[1]),
(pred_bbox[0] + pred_bbox[2], pred_bbox[1] + pred_bbox[3]),
(0, 255, 255), 3)
cv2.putText(img, str(idx), (40, 40), cv2.FONT_HERSHEY_SIMPLEX,
1, (0, 255, 255), 2)
cv2.imshow(video.name, img)
cv2.waitKey(1)
toc /= cv2.getTickFrequency()
# save results
model_path = os.path.join('results', args.dataset, model_name)
if not os.path.isdir(model_path):
os.makedirs(model_path)
result_path = os.path.join(model_path, '{}.txt'.format(video.name))
with open(result_path, 'w') as f:
for x in pred_bboxes:
f.write(','.join([str(i) for i in x]) + '\n')
print('({:3d}) Video: {:12s} Time: {:5.1f}s Speed: {:3.1f}fps'.format(
v_idx + 1, video.name, toc, idx / toc))
def main():
#try:
#os.remove("/home/developer/kashyap/pysot-master/*.csv")
#except:
# pass
# with open('./demo/groundtruth.csv', 'r') as f:
# reader = csv.reader(f)
# cords = list(reader)
# load config
cfg.merge_from_file('./experiments/siamrpn_alex_dwxcorr/config.yaml')
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
print(device)
# create model
model = ModelBuilder()
# load model
model.load_state_dict(torch.load('./experiments/siamrpn_alex_dwxcorr/model.pth', map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# build tracker
tracker = build_tracker(model)
video_list = glob1("/home/developer/kashyap/pysot-master/demo/vids/", "*.mp4")
for video_name in video_list:
video_name_str = os.path.splitext(video_name)[0]
df = pd.read_csv('./demo/vids/'+video_name_str+'.csv', delimiter=',', header=None)
cords = [list(x) for x in df.values]
object_counter = 0
for cord in cords:
object_counter = object_counter + 1
first_frame = True
# if video_name:#args.video_name:
# #video_name = args.video_name.split('/')[-1].split('.')[0]
# video_name = video_name.split('/')[-1].split('.')[0]
# else:
# exit()
frame_count = 1
mylist = [[frame_count,object_counter,cord,video_name]]
for frame in get_frames(video_name):#(args.video_name):
if first_frame:
try:
init_rect = cord
except:
exit()
tracker.init(frame, init_rect)
first_frame = False
else:
outputs = tracker.track(frame)
if 'polygon' in outputs:
exit()
else:
#crds = map(int,outputs['bbox'])
bbox = list(map(int,outputs['bbox']))
#cv2.rectangle(frame,(bbox[0],bbox[1]),(bbox[0]+bbox[2],bbox[1]+bbox[3]),(0,255,0),3)
#for frame in get_frames(video_name):#(args.video_name):
frame_count = frame_count + 1
mylist.append([frame_count,object_counter,bbox,video_name])
with open('vid-'+str(video_name)+'-tracking-'+str(object_counter)+'-object-'+str(cord)+'.csv', 'w', newline='') as csvfile:
writer = csv.writer(csvfile, quoting=0, '\n')#,quotechar='',escapechar='')
writer.writerow(mylist)
