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()