def __init__(self, count=5):
# init net
print('init tracker ...')
start = time()
self.net = SiamRPNvot()
self.net.load_state_dict(
torch.load(
'/home/car/projects/catkin_ws/src/tracking/scripts/DaSiamRPN/SiamRPNVOT.model'
))
self.net.eval().cuda()
self.yolo = detector()
self.human_detected = False
self.tracking = False
self.iou_threshold = 0.5
self.max_count = count
self.count = count
print('tracker initialize done, cost time {}'.format(time() - start))
def __init__(self, root="DaSiamRPN/", model="SiamRPNVOT.model", use_gpu=False):
"""
root(str): Persistent model data Directory.
model(str): name of the model[file] to run.
use_gpu(bool): enable gpu.
"""
# cache args
self.root_ = root
self.model_ = model
self.use_gpu_ = use_gpu
# Hard-coded sub-directory (TODO: maybe change this?)
model_path = join(realpath(dirname(__file__)), root, model)
# load net
self.net = SiamRPNvot()
self.net.load_state_dict(torch.load(model_path))
if self.use_gpu_:
self.net.eval().cuda()
else:
self.net.eval()
def main():
cap = cv2.VideoCapture(0)
net = SiamRPNvot()
net.load_state_dict(torch.load('./DaSiamRPN/SiamRPNVOT.model'))
net.eval().cuda()
_, im = cap.read()
init_box = cv2.selectROI('selectROI', im)
print(init_box)
target_pos, target_sz = rect_2_cxy_wh(init_box)
state = SiamRPN_init(im, target_pos, target_sz, net)
while True:
_, im = cap.read()
state = SiamRPN_track(state, im) # track
res = cxy_wh_2_rect(state['target_pos'], state['target_sz'])
res = [int(l) for l in res]
cv2.rectangle(im, (res[0], res[1]), (res[0] + res[2], res[1] + res[3]), (0, 255, 255), 3)
cv2.imshow('SiamRPN', im)
cv2.waitKey(1)
def main():
# load model
net = SiamRPNvot()
net.load_state_dict(torch.load('./DaSiamRPN/SiamRPNVOT.model'))
net.eval().cuda()
# image and init box
dataset_folder = './test_video/Panda'
ground_truth = np.loadtxt(join(dataset_folder, 'groundtruth_rect.txt'),
delimiter=',')
init_box = ground_truth[0, :]
image_list = [
join(dataset_folder, 'img', x)
for x in list(listdir(join(dataset_folder, 'img')))
]
# tracker init
target_pos, target_sz = rect_2_cxy_wh(init_box)
im = cv2.imread(image_list[0])
state = SiamRPN_init(im, target_pos, target_sz, net)
toc = 0
# tracking and visualization
for image_name in image_list[1:]:
im = cv2.imread(image_name)
tic = cv2.getTickCount()
state = SiamRPN_track(state, im) # track
toc += cv2.getTickCount() - tic
res = cxy_wh_2_rect(state['target_pos'], state['target_sz'])
res = [int(l) for l in res]
cv2.rectangle(im, (res[0], res[1]), (res[0] + res[2], res[1] + res[3]),
(0, 255, 255), 3)
cv2.imshow('SiamRPN', im)
cv2.waitKey(1)
print('Tracking Speed {:.1f}fps'.format(
(len(image_list) - 1) / (toc / cv2.getTickFrequency())))
class HumanTracker(object):
"""
Use DaSiamRPN and Yolov3 to track human
"""
def __init__(self, count=5):
# init net
print('init tracker ...')
start = time()
self.net = SiamRPNvot()
self.net.load_state_dict(
torch.load(
'/home/car/projects/catkin_ws/src/tracking/scripts/DaSiamRPN/SiamRPNVOT.model'
))
self.net.eval().cuda()
self.yolo = detector()
self.human_detected = False
self.tracking = False
self.iou_threshold = 0.5
self.max_count = count
self.count = count
print('tracker initialize done, cost time {}'.format(time() - start))
def track(self, img):
self.count -= 1
# get siamrpn tracker's result
state = SiamRPN_track(self.state, img)
res = cxy_wh_2_rect(state['target_pos'], state['target_sz'])
res = [int(l) for l in res]
if self.count <= 0:
# use detector to examine
human = self.yolo.detect_human(img)
if len(human) == 0:
self.human_detected = False
# if no detection, return sot result
return res
self.count = self.max_count
# compare iou between detect result and tracking result
bbox = [
human[0], human[1], human[2] - human[0], human[3] - human[1]
]
bbox = [int(l) for l in bbox]
if iou(bbox, res) < self.iou_threshold:
self.init_SOT(bbox, img)
return bbox
else:
return res
return res
def init_SOT(self, bbox, img):
start = time()
self.tracking = True
self.human_detected = True
target_pos, target_sz = rect_2_cxy_wh(bbox)
self.state = SiamRPN_init(img, target_pos, target_sz, self.net)
print('SOT init done, cost time: {}s'.format(time() - start))
def detect(self, img):
# detect human and return position
human = self.yolo.detect_human(img)
if len(human) == 0:
return (False, None)
else:
init_box = [
human[0], human[1], human[2] - human[0], human[3] - human[1]
]
return (True, init_box)
import torch
from os.path import join
from os import listdir
from DaSiamRPN.run_SiamRPN import SiamRPN_init, SiamRPN_track
from DaSiamRPN.net import SiamRPNvot
from DaSiamRPN.utils import cxy_wh_2_rect, rect_2_cxy_wh
from Yolov3.detector import detector
# init yolov3, video
yolo = detector()
detect_human = False
cap = cv2.VideoCapture(0)
# init tracking
net = SiamRPNvot()
net.load_state_dict(torch.load('./DaSiamRPN/SiamRPNVOT.model'))
net.eval().cuda()
init_box = []
while not detect_human:
_, img = cap.read()
human = yolo.detect_human(img)
print(human)
if len(human) == 0:
continue
else:
detect_human = True
init_box = [
human[0], human[1], human[2] - human[0], human[3] - human[1]
class Object_Tracker:
def __init__(self, root="DaSiamRPN/", model="SiamRPNVOT.model", use_gpu=False):
"""
root(str): Persistent model data Directory.
model(str): name of the model[file] to run.
use_gpu(bool): enable gpu.
"""
# cache args
self.root_ = root
self.model_ = model
self.use_gpu_ = use_gpu
# Hard-coded sub-directory (TODO: maybe change this?)
model_path = join(realpath(dirname(__file__)), root, model)
# load net
self.net = SiamRPNvot()
self.net.load_state_dict(torch.load(model_path))
if self.use_gpu_:
self.net.eval().cuda()
else:
self.net.eval()
@staticmethod
def get_confidence(state):
"""Returns confidence (score) of the current tracker output
Returns:
float: Score between 0 and 1 (where 1 is more confident)
"""
return state["score"]
def init(self, image, bounding_box):
"""Begins tracking of an object contained in a bounding box
Args:
image(A(H,W,3), uint8): Input image
bounding_box (4-tuple): Bounding box of object, in
form (center_x, center_y, width, height);
relative coordinates [0-1]
"""
if np.any(np.greater(np.asarray(bounding_box), 1.0)):
# input is probably not relative.
# we could resolve this, but we choose to throw an error.
raise ValueError("Please Make sure that the input box is relative !!")
img_h, img_w = np.shape(image)[:2]
[cx, cy, w, h] = utils.convert_to_pixels([img_w, img_h], bounding_box)
# tracker init
target_pos, target_sz = np.array([cx, cy]), np.array([w, h])
state = SiamRPN_init(
image, target_pos, target_sz, self.net, use_gpu=self.use_gpu_
)
return state
def __call__(self, image, box, state, display_tracking_speed=False):
"""Track the object in an image using the pre-calculated internal state
and a new image. (init must be run before this function.)
Args:
image(A(H,W,3), uint8): Input image
box : Unused; only here for compatibility.
display_tracking_speed (bool, optional): Option to print trakcing fps
Returns:
box(A(4)): Returns a relative bounding box of type [center_x, center_y, width, height]
state(?): Returns meta-information required for maintaining tracking state.
Note:
A(shape[,dtype=float]) in the above documentation is a shorthand for numpy.array.
"""
# tracking and visualization
toc = 0
tic = cv2.getTickCount()
state = SiamRPN_track(state, image, use_gpu=self.use_gpu_) # track
toc += cv2.getTickCount() - tic
box = (
state["target_pos"][0],
state["target_pos"][1],
state["target_sz"][0],
state["target_sz"][1],
)
img_h, img_w = np.shape(image)[:2]
box = utils.convert_to_relative([img_w, img_h], box)
return box, state