def __init__(self, root_dir, model_path, device):
self.root_dir = root_dir
self.device = device
self.transform = NormalizeToTensor()
self.scale = Rescale((224, 224))
self.model_path = model_path
self.model = GoNet()
self.gt = []
self.opts = None
self.curr_img = None
checkpoint = torch.load(
model_path, map_location=lambda storage, loc: storage)
self.model.load_state_dict(checkpoint['state_dict'])
self.model.to(device)
frames = [file for file in os.listdir(root_dir + '/img') if'jpeg' in file or 'jpg' in file]
frames = [root_dir + "/img/" + frame for frame in frames]
self.len = len(frames)-1
frames = np.array(frames)
frames.sort()
self.x = []
try:
f = open(root_dir + '/groundtruth_rect.txt')
except:
f = open(root_dir + '/groundtruth.txt')
lines = f.readlines()
#print('frame length is', self.len)
#print('gt length is', len(lines))
#choose min between gt and the frame for evaluation
self.len = min(self.len, len(lines)-1)
lines[0] = re.sub('\t', ',', lines[0])
lines[0] = re.sub(' +', ',', lines[0])
init_bbox = lines[0].strip().split(',')
init_bbox = [float(x) for x in init_bbox]
init_bbox = [init_bbox[0], init_bbox[1], init_bbox[0]+init_bbox[2],
init_bbox[1] + init_bbox[3]]
init_bbox = np.array(init_bbox)
self.prev_rect = init_bbox
self.img = []
for i in range(self.len):
self.x.append([frames[i], frames[i+1]])
img_prev = cv2.imread(frames[i])
img_prev = bgr2rgb(img_prev)
img_curr = cv2.imread(frames[i+1])
img_curr = bgr2rgb(img_curr)
self.img.append([img_prev, img_curr])
lines[i+1] = re.sub('\t', ',', lines[i+1])
lines[i+1] = re.sub(' +', ',', lines[i+1])
bb = lines[i+1].strip().split(',')
bb = [float(x) for x in bb]
bb = [bb[0], bb[1], bb[0]+bb[2], bb[1]+bb[3]]
self.gt.append(bb)
self.x = np.array(self.x)
print(init_bbox)
def _get_sample(self):
"""
Returns cropped previous and current frame at the previous predicted
location. Note that the images are scaled to (224,224,3).
"""
prev = self.prev_img
curr = self.curr_img
prevbb = self._last_bbox
prev_sample, opts_prev = crop_sample({'image': prev, 'bb': prevbb})
curr_sample, opts_curr = crop_sample({'image': curr, 'bb': prevbb})
prev_img = bgr2rgb(self.scale(prev_sample, opts_prev)['image'])
curr_img = bgr2rgb(self.scale(curr_sample, opts_curr)['image'])
sample = {'previmg': prev_img, 'currimg': curr_img}
self.curr_img = curr
self.opts = opts_curr
return sample
def get_orig_sample(self, idx):
"""
Returns original image with bounding box at a specific index.
Range of valid index: [0, self.len-1].
"""
curr = cv2.imread(self.x[idx])
curr = bgr2rgb(curr)
currbb = self.y[idx]
sample = {'image': curr, 'bb': currbb}
return sample
def __init__(self, root_dir, model_path, device, live=False):
self.root_dir = root_dir
self.device = device
self.transform = NormalizeToTensor()
self.scale = Rescale((224, 224))
self.model_path = model_path
self.model = GoNet()
self.gt = []
self.opts = None
self.curr_img = None
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
self.model.load_state_dict(checkpoint['state_dict'])
self.model.to(device)
if live:
# contains only img.
frames = glob.glob(os.path.join(root_dir, '*.jpg'))
frames = sorted(frames)
self.len = len(frames) - 1
frames = np.array(frames)
frames.sort()
self.x = []
self.img = []
for i in range(self.len):
self.x.append([frames[i], frames[i + 1]])
img_prev = cv2.imread(frames[i])
# img_prev = bgr2rgb(img_prev)
img_curr = cv2.imread(frames[i + 1])
# img_curr = bgr2rgb(img_curr)
self.img.append([img_prev, img_curr])
self.x = np.array(self.x)
else:
frames = os.listdir(root_dir + '/img')
frames = [root_dir + "/img/" + frame for frame in frames]
self.len = len(frames) - 1
frames = np.array(frames)
frames.sort()
self.x = []
f = open(root_dir + '/groundtruth_rect.txt')
lines = f.readlines()
lines[0] = re.sub('\t', ',', lines[0])
lines[0] = re.sub(' +', ',', lines[0])
init_bbox = lines[0].strip().split(',')
init_bbox = [float(x) for x in init_bbox]
init_bbox = [
init_bbox[0], init_bbox[1], init_bbox[0] + init_bbox[2],
init_bbox[1] + init_bbox[3]
]
init_bbox = np.array(init_bbox)
self.prev_rect = init_bbox
self.img = []
for i in range(self.len):
self.x.append([frames[i], frames[i + 1]])
img_prev = cv2.imread(frames[i])
img_prev = bgr2rgb(img_prev)
img_curr = cv2.imread(frames[i + 1])
img_curr = bgr2rgb(img_curr)
self.img.append([img_prev, img_curr])
lines[i + 1] = re.sub('\t', ',', lines[i + 1])
lines[i + 1] = re.sub(' +', ',', lines[i + 1])
bb = lines[i + 1].strip().split(',')
bb = [float(x) for x in bb]
bb = [bb[0], bb[1], bb[0] + bb[2], bb[1] + bb[3]]
self.gt.append(bb)
self.x = np.array(self.x)
print(init_bbox)
