Function run_each_dataset

tutorials/ctracker/test.py:135–281

(model_dir, retinanet, dataset_path, subset, cur_dataset)

def run_each_dataset(model_dir, retinanet, dataset_path, subset, cur_dataset):	
	print(cur_dataset)

	img_list = os.listdir(os.path.join(dataset_path, subset, cur_dataset, 'img1'))
	img_list = [os.path.join(dataset_path, subset, cur_dataset, 'img1', _) for _ in img_list if ('jpg' in _) or ('png' in _)]
	img_list = sorted(img_list)

	img_len = len(img_list)
	last_feat = None

	confidence_threshold = 0.4
	IOU_threshold = 0.5
	retention_threshold = 10

	det_list_all = []
	tracklet_all = []
	max_id = 0
	max_draw_len = 100
	draw_interval = 5
	img_width = 1920
	img_height = 1080
	fps = 30

	for i in range(img_len):
		det_list_all.append([])

	for idx in range((int(img_len / 2)), img_len + 1):
		i = idx - 1
		print('tracking: ', i)
		with torch.no_grad():
			data_path1 = img_list[min(idx, img_len - 1)]
			img_origin1 = skimage.io.imread(data_path1)
			img_h, img_w, _ = img_origin1.shape
			img_height, img_width = img_h, img_w
			resize_h, resize_w = math.ceil(img_h / 32) * 32, math.ceil(img_w / 32) * 32
			img1 = np.zeros((resize_h, resize_w, 3), dtype=img_origin1.dtype)
			img1[:img_h, :img_w, :] = img_origin1
			img1 = (img1.astype(np.float32) / 255.0 - np.array([[RGB_MEAN]])) / np.array([[RGB_STD]])
			img1 = torch.from_numpy(img1).permute(2, 0, 1).view(1, 3, resize_h, resize_w)
			scores, transformed_anchors, last_feat = retinanet(img1.cuda().float(), last_feat=last_feat)
# 			if idx > 0:
			if idx > (int(img_len / 2)):
				idxs = np.where(scores>0.1)

				for j in range(idxs[0].shape[0]):
					bbox = transformed_anchors[idxs[0][j], :]
					x1 = int(bbox[0])
					y1 = int(bbox[1])
					x2 = int(bbox[2])
					y2 = int(bbox[3])

					x3 = int(bbox[4])
					y3 = int(bbox[5])
					x4 = int(bbox[6])
					y4 = int(bbox[7])

					det_conf = float(scores[idxs[0][j]])