[Suggest]: The traversal method used to create the mask in visualize_token_selection.py is inefficient.

[tiantianNewer]

LAST-ViT/visualization/visualize_token_selection.py

Lines 475 to 478 in cdeb884

	mask = np.zeros(num_tokens)
	for token_idx in range(num_tokens):
	if token_freq[token_idx] > threshold:
	mask[token_idx] = 1

LAST-ViT/visualization/visualize_token_selection.py

Lines 588 to 591 in cdeb884

	mask = np.zeros(num_tokens)
	for token_idx in range(num_tokens):
	if token_freq[token_idx] > threshold:
	mask[token_idx] = 1

suggest changing to:

mask = token_freq > threshold

[tiantianNewer]

LAST-ViT/visualization/visualize_token_selection.py

Lines 361 to 373 in cdeb884

	# Expand mask to image size
	mask_img = np.zeros((image_size, image_size), dtype=bool)
	patch_h = image_size // grid_size
	patch_w = image_size // grid_size
	for i in range(grid_size):
	for j in range(grid_size):
	if mask_2d[i, j] > 0:
	h_start = i * patch_h
	h_end = min((i + 1) * patch_h, image_size)
	w_start = j * patch_w
	w_end = min((j + 1) * patch_w, image_size)
	mask_img[h_start:h_end, w_start:w_end] = True

here also suggest changing to:

# Expand mask to image size
    patch_h = image_size // grid_size
    patch_w = image_size // grid_size
    mask_img = np.kron(mask_2d, np.ones((patch_h, patch_w), dtype=bool))