class SiamRPNPP(BaseTracker):
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 initialize(self, image_file, box):
image = cv2.imread(image_file)
self.tracker.init(image, box)
def track(self, image_file):
image = cv2.imread(image_file)
bbox = self.tracker.track(image)["bbox"]
return bbox
class SiamRPNPP():
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 get_heat_map(self, X_crop, softmax=False):
score_map = self.model.track(X_crop)['cls']#(N,2x5,25,25)
score_map = score_map.permute(1, 2, 3, 0).contiguous().view(2, -1).permute(1, 0)#(5HWN,2)
if softmax:
score_map = F.softmax(score_map, dim=1).data[:, 1]#(5HWN,)
return score_map
def get_cls_reg(self, X_crop, softmax=False):
outputs = self.model.track(X_crop)#(N,2x5,25,25)
score_map = outputs['cls'].permute(1, 2, 3, 0).contiguous().view(2, -1).permute(1, 0)#(5HWN,2)
reg_res = outputs['loc'].permute(1, 2, 3, 0).contiguous().view(4, -1)
if softmax:
score_map = F.softmax(score_map, dim=1).data[:, 1]#(5HWN,)
return score_map, reg_res
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
class SiamRPNPPGroup(BaseTracker):
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 initialize(self, image_file, box):
image = cv2.imread(image_file)
self.tracker.init(image, box)
def track(self, image_file):
image = cv2.imread(image_file)
bbox = self.tracker.track(image)["bbox"]
return bbox
class Tracking(object):
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, frame, init_rect):
print("initial rectangle selected as: ", init_rect)
init_rect = list(map(int, init_rect.split(',')))
self.tracker.init(frame, init_rect)
def update(self, bbox):
## REMEMBER TO CALL UPDATE
self.tracker.update(bbox)
def get_roi(self, img, instance_size):
return self.tracker.get_roi(img, instance_size)
def track(self, frame, x_crop, scale_z, instance_size):
# x_crop, scale_z = self.get_roi(frame)
return self.tracker.track(frame, x_crop, scale_z, instance_size)
# Following functions are used for template update
def templateFeature(self, z):
return self.model.templateFeature(z)
def zf(self):
return self.model.zf
def updateTemplate(self, zf):
model.zf = zf
def generate_data():
# Load config
cfg.merge_from_file("experiments/siammask_r50_l3/config.yaml")
with torch.no_grad():
# Load forward model
model = ModelBuilder()
model.load_state_dict(
torch.load("experiments/siammask_r50_l3/model.pth",
map_location=CPU))
model.share_memory()
model.eval().to(DEVICE)
detection_by_tracking(frame_dir=args.frame_dir,
json_file=args.json_file,
detection_threshold=0.8,
tracking_threshold=0.9,
save_json_file=args.tracking_json,
tracker_model=model,
offset=0,
high=12,
low=12,
step=24,
parallel=False,
multithreading=False)
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 main():
# load config
print("begin")
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_root = "/data/VisDrone Challenge/Single-Object Tracking/VisDrone2019-SOT-val/"
# create model
model = ModelBuilder()
tracker = build_tracker(model)
dataset = DatasetFactory.create_dataset(name=args.dataset,
dataset_root=dataset_root,
load_img=False)
print(dataset.dataset_root)
j =0
for v_idx, video in enumerate(dataset):
for idx, (img, gt_bbox) in enumerate(video):
if idx == 0:
cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_bbox))
# 左上角坐标 ，w ，h 形式
gt_bbox_ = [cx - (w - 1) / 2, cy - (h - 1) / 2, w, h]
# 初始化tracker 的box
img=tracker.init(img, gt_bbox_)
img =img.cpu().numpy()[0].transpose(1,2,0)
cv2.imwrite("./Radio{:06d}.jpg".format(j),img)
j += 1
else:
break
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 save_siamese_rpn():
# load config
rpn_path = root_dir + 'experiments/siamrpn_alex_dwxcorr_16gpu/pre_train/checkpoint_e45.pth'
gru_rpn = root_dir + 'experiments/siamrpn_alex_dwxcorr_16gpu/config.yaml'
cfg.merge_from_file(gru_rpn)
# create model
model_rpn = ModelBuilder()
model_rpn = load_pretrain(model_rpn, rpn_path).cuda().eval()
gru_path = root_dir + 'experiments/siamrpn_alex_dwxcorr_16gpu/gru_snapshot/gru_10.pth'
gru_cfg = root_dir + 'experiments/siamrpn_alex_dwxcorr_16gpu/config_gru.yaml'
cfg.merge_from_file(gru_cfg)
# create model
model_gru = ModelBuilder()
model_gru = load_pretrain(model_gru, gru_path).cuda().eval()
for key, item in model_gru.named_parameters():
# print(key.find("grus"))
print(key, item.shape)
for key, item in model_rpn.named_parameters():
# print(key.find("grus"))
print(key, item.shape)
model_gru_dict = model_gru.state_dict()
model_rpn_dict = model_rpn.state_dict()
for key in model_gru_dict:
if key.find("grus") != -1:
print("fix:", key)
else:
print("change:", key)
model_gru_dict[key] = model_rpn_dict[key]
# 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_gru_dict, "siamese_gru10_rpn45.pth")
def main():
rank, world_size = dist_init()
# rank, world_size = 0,1
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(cfg)
# 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)
# load pretrain
elif cfg.TRAIN.PRETRAINED:
load_pretrain(model, cfg.TRAIN.PRETRAINED)
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 __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 setup_tracker():
cfg.merge_from_file(cfg_file)
model = ModelBuilder()
model = load_pretrain(model, model_file).cuda().eval()
tracker = build_tracker(model)
warmup(model)
return tracker
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 __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():
#############################################################################
# initialize the tracker
cfg.merge_from_file("experiments/siamrpn_mobilev2_l234_dwxcorr/config.yaml")
cfg.CUDA = torch.cuda.is_available()
device = torch.device("cuda" if cfg.CUDA else "cpu")
model = ModelBuilder()
model.load_state_dict(
torch.load("experiments/siamrpn_mobilev2_l234_dwxcorr/model.pth",
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
tracker = build_tracker(model)
#############################################################################
# initialzie the benchmark parameter
img = cv2.imread("image/benchmark_5.jpg")
bbox = (131, 122, 92, 118)
# calculate channle average
channel_average = np.mean(img, axis=(0, 1))
# EXEMPLAR_SIZE of mobilenetV2 is 127
z_new_crop = get_subwindow(img, 127, bbox, channel_average)
print(z_new_crop.mean(), "\t", z_new_crop.std())
def __init__(self,classid=0,anchorid=0,init_frame=None,init_bbox=None):
# load config
cfg_path = '../object_detection/pysot/experiments/siamrpn_r50_l234_dwxcorr/config.yaml'
snapshot = '../object_detection/pysot/experiments/siamrpn_r50_l234_dwxcorr/model.pth'
cfg.merge_from_file(cfg_path)
cfg.CUDA = torch.cuda.is_available() and cfg.CUDA
# cfg.CUDA = False
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# device='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)
self.tracker = tracker
self.classId = classid
self.anchorId = anchorid
self.init_bbox = init_bbox
self.init_frame = init_frame
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_state_dict(
torch.load(
args.snapshot,
map_location=lambda storage, loc: storage.cpu())['state_dict'])
model.eval().to(device)
# build tracker
if cfg.RPN.TYPE == "YOLO":
tracker = build_tracker(model, True)
else:
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)
writer = cv2.VideoWriter("result.avi", cv2.VideoWriter_fourcc(*'XVID'),
25.0, (480, 360))
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)
writer.write(frame)
cv2.waitKey(40)
writer.release()
def __init__(self, parent=None):
super(MyMainWindow, self).__init__(parent)
# 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.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
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())['state_dict'])
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)
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 cfg.TRANSFORMER.TRANSFORMER:
acc, (x1, y1, x2, y2) = outputs
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 3)
cv2.putText(frame, 'Acc: ' + acc.astype('str'), (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
cv2.imshow(video_name, frame)
cv2.waitKey(40)
else:
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)
def __init__(self, config_file, model_path):
super().__init__()
cfg.merge_from_file(config_file)
model = ModelBuilder()
model.load_state_dict(
torch.load(model_path,
map_location=lambda storage, loc: storage.cpu()))
model.eval().cuda()
self.tracker = build_tracker(model)
def build_model(self):
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)
return tracker
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.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 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 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 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 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 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()
