init

visual_search.py

@@ -0,0 +1,567 @@
+import argparse
+import os
+import sys
+import json
+import tqdm
+import copy
+from queue import PriorityQueue
+import functools
+import spacy
+nlp = spacy.load("en_core_web_sm")
+
+import cv2
+from PIL import Image
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, CLIPImageProcessor
+from transformers import OwlViTProcessor
+
+from VisualSearch.model.VSM import VSMForCausalLM
+from VisualSearch.model.llava import conversation as conversation_lib
+from VisualSearch.model.llava.mm_utils import tokenizer_image_token
+from VisualSearch.utils.utils import expand2square
+from VisualSearch.utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
+						 DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX)
+
+
+def parse_args(args):
+	parser = argparse.ArgumentParser(description="Visual Search Evaluation")
+	parser.add_argument("--version", default="craigwu/seal_vsm_7b")
+	parser.add_argument("--benchmark-folder", default="vstar_bench", type=str)
+	parser.add_argument("--visualization", action="store_true", default=False)
+	parser.add_argument("--output_path", default="", type=str)
+	parser.add_argument("--confidence_low", default=0.3, type=float)
+	parser.add_argument("--confidence_high", default=0.5, type=float)
+	parser.add_argument("--target_cue_threshold", default=6.0, type=float)
+	parser.add_argument("--target_cue_threshold_decay", default=0.7, type=float)
+	parser.add_argument("--target_cue_threshold_minimum", default=3.0, type=float)
+	parser.add_argument("--minimum_size_scale", default=4.0, type=float)
+	parser.add_argument("--minimum_size", default=224, type=int)
+	parser.add_argument("--model_max_length", default=512, type=int)
+	parser.add_argument(
+		"--vision-tower", default="openai/clip-vit-large-patch14", type=str
+	)
+	parser.add_argument("--use_mm_start_end", action="store_true", default=True)
+	parser.add_argument(
+		"--conv_type",
+		default="llava_v1",
+		type=str,
+		choices=["llava_v1", "llava_llama_2"],
+	)
+	return parser.parse_args(args)
+
+def tranverse(token):
+	children = [_ for _ in token.children]
+	if len(children) == 0:
+		return token.i, token.i
+	left_i = token.i
+	right_i = token.i
+	for child in children:
+		child_left_i, child_right_i = tranverse(child)
+		left_i = min(left_i, child_left_i)
+		right_i = max(right_i, child_right_i)
+	return left_i, right_i
+def get_noun_chunks(token):
+	left_children = []
+	right_children = []
+	for child in token.children:
+		if child.i < token.i:
+			left_children.append(child)
+		else:
+			right_children.append(child)
+
+	start_token_i = token.i
+	for left_child in left_children[::-1]:
+		if left_child.dep_ in ['amod', 'compound', 'poss']:
+			start_token_i, _ = tranverse(left_child)
+		else:
+			break
+	end_token_i = token.i
+	for right_child in right_children:
+		if right_child.dep_ in ['relcl', 'prep']:
+			_, end_token_i = tranverse(right_child)
+		else:
+			break
+	return start_token_i, end_token_i
+
+def filter_chunk_list(chunks):
+	def overlap(min1, max1, min2, max2):
+		return min(max1, max2) - max(min1, min2)
+	chunks = sorted(chunks, key=lambda chunk: chunk[1]-chunk[0], reverse=True)
+	filtered_chunks = []
+	for chunk in chunks:
+		flag=True
+		for exist_chunk in filtered_chunks:
+			if overlap(exist_chunk[0], exist_chunk[1], chunk[0], chunk[1]) >= 0:
+				flag = False
+				break
+		if flag:
+			filtered_chunks.append(chunk)
+	return sorted(filtered_chunks, key=lambda chunk: chunk[0])
+
+def extract_noun_chunks(expression):
+	doc = nlp(expression)
+	cur_chunks = []
+	for token in doc:
+		if token.pos_ not in ["NOUN", "PRON"]:
+			continue
+		cur_chunks.append(get_noun_chunks(token))
+	cur_chunks = filter_chunk_list(cur_chunks)
+	cur_chunks = [doc[chunk[0]:chunk[1]+1].text for chunk in cur_chunks]
+	return cur_chunks
+
+def preprocess(
+	x,
+	pixel_mean=torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1),
+	pixel_std=torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1),
+	img_size=1024,
+) -> torch.Tensor:
+	"""Normalize pixel values and pad to a square input."""
+	# Normalize colors
+	x = (x - pixel_mean) / pixel_std
+	# Pad
+	h, w = x.shape[-2:]
+	padh = img_size - h
+	padw = img_size - w
+	x = F.pad(x, (0, padw, 0, padh))
+	return x
+
+def box_cxcywh_to_xyxy(x):
+	x_c, y_c, w, h = x.unbind(1)
+	b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
+		 (x_c + 0.5 * w), (y_c + 0.5 * h)]
+	return torch.stack(b, dim=1)
+
+def rescale_bboxes(out_bbox, size):
+	img_w, img_h = size
+	b = box_cxcywh_to_xyxy(out_bbox)
+	b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
+	return b
+
+class VSM:
+	def __init__(self, args):
+		kwargs = {}
+		kwargs['torch_dtype'] = torch.bfloat16
+		kwargs['device_map'] = 'cuda'
+		kwargs['is_eval'] = True
+		vsm_tokenizer = AutoTokenizer.from_pretrained(
+				args.version,
+				cache_dir=None,
+				model_max_length=args.model_max_length,
+				padding_side="right",
+				use_fast=False,
+			)
+		vsm_tokenizer.pad_token = vsm_tokenizer.unk_token
+		loc_token_idx = vsm_tokenizer("[LOC]", add_special_tokens=False).input_ids[0]
+		vsm_model = VSMForCausalLM.from_pretrained(
+				args.version, low_cpu_mem_usage=True, vision_tower=args.vision_tower, loc_token_idx=loc_token_idx, **kwargs
+			)
+		vsm_model.get_model().initialize_vision_modules(vsm_model.get_model().config)
+		vision_tower = vsm_model.get_model().get_vision_tower().cuda().to(dtype=torch.bfloat16)
+		vsm_image_processor = vision_tower.image_processor
+		vsm_model.eval()
+		clip_image_processor = CLIPImageProcessor.from_pretrained(vsm_model.config.vision_tower)
+		transform = OwlViTProcessor.from_pretrained("google/owlvit-base-patch16")
+		self.model = vsm_model
+		self.vsm_tokenizer = vsm_tokenizer
+		self.vsm_image_processor = vsm_image_processor
+		self.clip_image_processor = clip_image_processor
+		self.transform = transform
+		self.conv_type = args.conv_type
+		self.use_mm_start_end = args.use_mm_start_end
+	
+	@torch.inference_mode()
+	def inference(self, image, question, mode='segmentation'):
+		conv = conversation_lib.conv_templates[self.conv_type].copy()
+		conv.messages = []
+		prompt = DEFAULT_IMAGE_TOKEN + "\n" + question
+		if self.use_mm_start_end:
+			replace_token = ( DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN)
+			prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+		conv.append_message(conv.roles[0], prompt)
+		conv.append_message(conv.roles[1], "")
+		prompt = conv.get_prompt()
+
+		background_color = tuple(int(x*255) for x in self.clip_image_processor.image_mean)
+		image_clip = self.clip_image_processor.preprocess(expand2square(image, background_color), return_tensors="pt")["pixel_values"][0].unsqueeze(0).cuda()
+
+		image_clip = image_clip.bfloat16()
+		image = np.array(image)
+		original_size_list = [image.shape[:2]]
+		image = self.transform(images=image, return_tensors="pt")['pixel_values'].cuda()
+		resize_list = [image.shape[:2]]
+		image = image.bfloat16()
+		input_ids = tokenizer_image_token(prompt, self.vsm_tokenizer, return_tensors="pt")
+		input_ids = input_ids.unsqueeze(0).cuda()
+
+		output_ids, pred_masks, det_result = self.model.inference(
+			image_clip,
+			image,
+			input_ids,
+			resize_list,
+			original_size_list,
+			max_new_tokens=100,
+			tokenizer=self.vsm_tokenizer,
+			mode = mode
+		)
+		if mode == 'segmentation':
+			pred_mask = pred_masks[0]
+			pred_mask = torch.clamp(pred_mask, min=0)
+			return pred_mask[-1]
+
+		elif mode == 'vqa':
+			input_token_len = input_ids.shape[1]
+			n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
+			if n_diff_input_output > 0:
+				print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+			text_output = self.vsm_tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
+			text_output = text_output.replace("\n", "").replace("  ", " ").strip()
+			return text_output
+		
+		elif mode == 'detection':
+			pred_mask = pred_masks[0]
+			pred_mask = torch.clamp(pred_mask, min=0)
+			return det_result['pred_boxes'][0].cpu(), det_result['pred_logits'][0].sigmoid().cpu(), pred_mask[-1]
+
+def refine_bbox(bbox, image_width, image_height):
+	bbox[0] = max(0, bbox[0])
+	bbox[1] = max(0, bbox[1])
+	bbox[2] = min(bbox[2], image_width-bbox[0])
+	bbox[3] = min(bbox[3], image_height-bbox[1])
+	return bbox
+
+def split_4subpatches(current_patch_bbox):
+	hw_ratio = current_patch_bbox[3] / current_patch_bbox[2]
+	if hw_ratio >= 2:
+		return 1, 4
+	elif hw_ratio <= 0.5:
+		return 4, 1
+	else:
+		return 2, 2
+
+def get_sub_patches(current_patch_bbox, num_of_width_patches, num_of_height_patches):
+	width_stride = int(current_patch_bbox[2]//num_of_width_patches)
+	height_stride = int(current_patch_bbox[3]/num_of_height_patches)
+	sub_patches = []
+	for j in range(num_of_height_patches):
+		for i in range(num_of_width_patches):
+			sub_patch_width = current_patch_bbox[2] - i*width_stride if i == num_of_width_patches-1 else width_stride
+			sub_patch_height = current_patch_bbox[3] - j*height_stride if j == num_of_height_patches-1 else height_stride
+			sub_patch = [current_patch_bbox[0]+i*width_stride, current_patch_bbox[1]+j*height_stride, sub_patch_width, sub_patch_height]
+			sub_patches.append(sub_patch)
+	return sub_patches, width_stride, height_stride
+
+def get_subpatch_scores(score_heatmap, current_patch_bbox, sub_patches):
+	total_sum = (score_heatmap/(current_patch_bbox[2]*current_patch_bbox[3])).sum()
+	sub_scores = []
+	for sub_patch in sub_patches:
+		bbox = [(sub_patch[0]-current_patch_bbox[0]), sub_patch[1]-current_patch_bbox[1], sub_patch[2], sub_patch[3]]
+		score = (score_heatmap[bbox[1]:bbox[1]+bbox[3], bbox[0]:bbox[0]+bbox[2]]/(current_patch_bbox[2]*current_patch_bbox[3])).sum()
+		if total_sum > 0:
+			score /= total_sum
+		else:
+			score *= 0
+		sub_scores.append(score)
+	return sub_scores
+
+def normalize_score(score_heatmap):
+	max_score = score_heatmap.max()
+	min_score = score_heatmap.min()
+	if max_score != min_score:
+		score_heatmap = (score_heatmap - min_score) / (max_score - min_score)
+	else:
+		score_heatmap = score_heatmap * 0
+	return score_heatmap
+
+def iou(bbox1, bbox2):
+	x1 = max(bbox1[0], bbox2[0])
+	y1 = max(bbox1[1], bbox2[1])
+	x2 = min(bbox1[0]+bbox1[2], bbox2[0]+bbox2[2])
+	y2 = min(bbox1[1]+bbox1[3],bbox2[1]+bbox2[3])
+	inter_area = max(0, x2 - x1) * max(0, y2 - y1)
+	return inter_area/(bbox1[2]*bbox1[3]+bbox2[2]*bbox2[3]-inter_area)
+
+BOX_COLOR = (255, 0, 0) # Red
+TEXT_COLOR = (255, 255, 255) # White
+import cv2
+from matplotlib import pyplot as plt
+def visualize_bbox(img, bbox, class_name, color=BOX_COLOR, thickness=2):
+	"""Visualizes a single bounding box on the image"""
+	x_min, y_min, w, h = bbox
+	x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)
+   
+	cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)
+	
+	((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)    
+	cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
+	cv2.putText(
+		img,
+		text=class_name,
+		org=(x_min, y_min - int(0.3 * text_height)),
+		fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+		fontScale=0.5, 
+		color=TEXT_COLOR, 
+		lineType=cv2.LINE_AA,
+	)
+	return img
+def show_heatmap_on_image(img: np.ndarray,
+					  mask: np.ndarray,
+					  use_rgb: bool = False,
+					  colormap: int = cv2.COLORMAP_JET,
+					  image_weight: float = 0.5) -> np.ndarray:
+	mask = np.clip(mask, 0, 1)
+	heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)
+	if use_rgb:
+		heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+	heatmap = np.float32(heatmap) / 255
+
+	if np.max(img) > 1:
+		raise Exception(
+			"The input image should np.float32 in the range [0, 1]")
+
+	if image_weight < 0 or image_weight > 1:
+		raise Exception(
+			f"image_weight should be in the range [0, 1].\
+				Got: {image_weight}")
+
+	cam = (1 - image_weight) * heatmap + image_weight * img
+	cam = cam / np.max(cam)
+	return np.uint8(255 * cam)
+def vis_heatmap(image, heatmap, use_rgb=False):
+	max_v =  np.max(heatmap)
+	min_v =  np.min(heatmap)
+	if max_v != min_v:
+		heatmap = (heatmap - min_v) / (max_v - min_v)
+	heatmap_image = show_heatmap_on_image(image.astype(float)/255., heatmap, use_rgb=use_rgb)
+	return heatmap_image
+
+def visualize_search_path(image, search_path, search_length, target_bbox, label, save_path):
+	context_cue_list = []
+	whole_image = image	
+	os.makedirs(save_path, exist_ok=True)
+	whole_image.save(os.path.join(save_path, 'whole_image.jpg'))
+
+	whole_image = np.array(whole_image)
+	if target_bbox is not None:
+		whole_image = visualize_bbox(whole_image.copy(), target_bbox, class_name="gt: "+label, color=(255,0,0))
+	for step_i, node in enumerate(search_path):
+		if step_i + 1 > search_length:
+			break
+		current_patch_box = node['bbox']
+		if 'detection_result' in node:
+			final_patch_image = image.crop((current_patch_box[0],current_patch_box[1],current_patch_box[0]+current_patch_box[2], current_patch_box[1]+current_patch_box[3]))
+			final_patch_image.save(os.path.join(save_path, 'final_patch_image.jpg'))
+			final_search_result = visualize_bbox(np.array(final_patch_image), node['detection_result'], class_name='search result', color=(255,0,0))
+			final_search_result = cv2.cvtColor(final_search_result, cv2.COLOR_RGB2BGR)
+			cv2.imwrite(os.path.join(save_path, 'search_result.jpg'), final_search_result)
+		cur_whole_image = visualize_bbox(whole_image.copy(), current_patch_box, class_name="step-{}".format(step_i+1), color=(0,0,255))
+		# if step_i != len(search_path)-1:
+		# 	next_patch_box = search_path[step_i+1]['bbox']
+		# 	cur_whole_image = visualize_bbox(cur_whole_image, next_patch_box, class_name="next-step", color=(0,255,0))
+		cur_whole_image = cv2.cvtColor(cur_whole_image, cv2.COLOR_RGB2BGR)
+		cv2.imwrite(os.path.join(save_path, 'step_{}.jpg'.format(step_i+1)), cur_whole_image)
+		
+		cur_patch_image = image.crop((current_patch_box[0],current_patch_box[1],current_patch_box[0]+current_patch_box[2], current_patch_box[1]+current_patch_box[3]))
+		if 'context_cue' in node:
+			context_cue = node['context_cue']
+			context_cue_list.append('step{}: {}'.format(step_i+1, context_cue)+'\n')
+		if 'final_heatmap' in node:
+			score_map = node['final_heatmap']
+			score_map = vis_heatmap(np.array(cur_patch_image), score_map, use_rgb=True)
+			score_map = cv2.cvtColor(score_map, cv2.COLOR_RGB2BGR)
+			cv2.imwrite(os.path.join(save_path, 'step_{}_heatmap.jpg'.format(step_i+1)), score_map)
+
+	with open(os.path.join(save_path, 'context_cue.txt'),"w") as f:
+		f.writelines(context_cue_list)
+		
+@functools.total_ordering
+class Prioritize:
+
+	def __init__(self, priority, item):
+		self.priority = priority
+		self.item = item
+
+	def __eq__(self, other):
+		return self.priority == other.priority
+
+	def __lt__(self, other):
+		return self.priority < other.priority
+def visual_search_queue(vsm, image, target_object_name, current_patch, search_path, queue,  smallest_size=224, confidence_high=0.5, target_cue_threshold=6.0, target_cue_threshold_decay=0.7, target_cue_threshold_minimum=3.0):
+	current_patch_bbox = current_patch['bbox']
+	current_patch_scale_level = current_patch['scale_level']
+
+	image_patch = image.crop((int(current_patch_bbox[0]), int(current_patch_bbox[1]), int(current_patch_bbox[0]+current_patch_bbox[2]), int(current_patch_bbox[1]+current_patch_bbox[3])))
+	# whehter we can detect the target object on the current image patch
+	question = "Please locate the {} in this image.".format(target_object_name)
+	pred_bboxes, pred_logits, target_cue_heatmap = vsm.inference(copy.deepcopy(image_patch), question, mode='detection')
+	if len(pred_logits) > 0:
+		top_index = pred_logits.view(-1).argmax()
+		top_logit = pred_logits.view(-1).max()
+		final_bbox = pred_bboxes[top_index].view(4)
+		final_bbox = final_bbox * torch.Tensor([image_patch.width, image_patch.height, image_patch.width, image_patch.height])
+		final_bbox[:2] -= final_bbox[2:] / 2
+		if top_logit > confidence_high:
+			search_path[-1]['detection_result'] = final_bbox
+			# only return multiple detected instances on the whole image
+			if len(search_path) == 1:
+				all_valid_boxes = pred_bboxes[pred_logits.view(-1)>0.5].view(-1, 4)
+				all_valid_boxes = all_valid_boxes * torch.Tensor([[image_patch.width, image_patch.height, image_patch.width, image_patch.height]])
+				all_valid_boxes[:, :2] -= all_valid_boxes[:, 2:] / 2
+				return True, search_path, all_valid_boxes
+			return True, search_path, None
+		else:
+			search_path[-1]['temp_detection_result'] = (top_logit, final_bbox)
+
+	### current patch is already the smallest unit
+	if min(current_patch_bbox[2], current_patch_bbox[3]) <= smallest_size:
+		return False, search_path, None
+
+	target_cue_heatmap = target_cue_heatmap.view(current_patch_bbox[3], current_patch_bbox[2], 1)
+	score_max = target_cue_heatmap.max().item()
+	# check whether the target cue is prominent
+	threshold = max(target_cue_threshold_minimum, target_cue_threshold*(target_cue_threshold_decay)**(current_patch_scale_level-1))
+	if score_max > threshold:
+		target_cue_heatmap = normalize_score(target_cue_heatmap)
+		final_heatmap = target_cue_heatmap
+	else:
+		question = "According to the common sense knowledge and possible visual cues, what is the most likely location of the {} in the image?".format(target_object_name)
+		vqa_results = vsm.inference(copy.deepcopy(image_patch), question, mode='vqa')
+
+		possible_location_phrase = vqa_results.split('most likely to appear')[-1].strip()
+		if possible_location_phrase.endswith('.'):
+			possible_location_phrase = possible_location_phrase[:-1]
+		possible_location_phrase = possible_location_phrase.split(target_object_name)[-1]
+		noun_chunks = extract_noun_chunks(possible_location_phrase)
+		if len(noun_chunks) == 1:
+			possible_location_phrase = noun_chunks[0]
+		else:
+			possible_location_phrase = "region {}".format(possible_location_phrase)
+		question = "Please locate the {} in this image.".format(possible_location_phrase)
+		context_cue_heatmap = vsm.inference(copy.deepcopy(image_patch), question, mode='segmentation').view(current_patch_bbox[3], current_patch_bbox[2], 1)
+		context_cue_heatmap = normalize_score(context_cue_heatmap)
+		final_heatmap = context_cue_heatmap
+
+	current_patch_index = len(search_path)-1
+	if score_max <= threshold:
+		search_path[current_patch_index]['context_cue'] = vqa_results + "#" + possible_location_phrase
+	search_path[current_patch_index]['final_heatmap'] = final_heatmap.cpu().numpy()
+	
+	### split the current patch into 4 sub-patches
+	basic_sub_patches, sub_patch_width, sub_patch_height = get_sub_patches(current_patch_bbox, *split_4subpatches(current_patch_bbox))
+
+	tmp_patch = current_patch
+	basic_sub_scores = [0]*len(basic_sub_patches)
+	while True:
+		tmp_score_heatmap = tmp_patch['final_heatmap']
+		tmp_sub_scores = get_subpatch_scores(tmp_score_heatmap, tmp_patch['bbox'],  basic_sub_patches)
+		basic_sub_scores = [basic_sub_scores[patch_i]+tmp_sub_scores[patch_i]/(4**tmp_patch['scale_level']) for patch_i in range(len(basic_sub_scores))]
+		if  tmp_patch['parent_index'] == -1:
+			break
+		else:
+			tmp_patch = search_path[tmp_patch['parent_index']]
+
+	sub_patches = basic_sub_patches
+	sub_scores = basic_sub_scores
+
+	for sub_patch, sub_score in zip(sub_patches, sub_scores):
+		new_patch_info = dict()
+		new_patch_info['bbox'] = sub_patch
+		new_patch_info['scale_level'] = current_patch_scale_level + 1
+		new_patch_info['score'] = sub_score
+		new_patch_info['parent_index'] = current_patch_index
+		queue.put(Prioritize(-new_patch_info['score'], new_patch_info))
+	
+	while(not queue.empty()):
+		patch_chosen = queue.get().item
+		search_path.append(patch_chosen)
+		success, search_path, all_valid_boxes = visual_search_queue(vsm, image, target_object_name, patch_chosen, search_path, queue, smallest_size=smallest_size, confidence_high=confidence_high, target_cue_threshold=target_cue_threshold, target_cue_threshold_decay=target_cue_threshold_decay, target_cue_threshold_minimum=target_cue_threshold_minimum)
+		if success:
+			return success, search_path, all_valid_boxes
+	return False, search_path, None
+
+
+def visual_search(vsm, image, target_object_name, target_bbox, smallest_size, confidence_high=0.5, confidence_low=0.3, target_cue_threshold=6.0, target_cue_threshold_decay=0.7, target_cue_threshold_minimum=3.0, visualize=False, save_path=None):
+	if visualize:
+		assert save_path is not None
+	init_patch = dict()
+	init_patch['bbox'] = [0,0,image.width,image.height]
+	init_patch['scale_level'] = 1
+	init_patch['score'] = None
+	init_patch['parent_index'] = -1
+	search_path = [init_patch]
+
+	queue = PriorityQueue()
+	search_successful, search_path, all_valid_boxes = visual_search_queue(vsm, image, target_object_name, init_patch, search_path, queue, smallest_size=smallest_size, confidence_high=confidence_high, target_cue_threshold=target_cue_threshold, target_cue_threshold_decay=target_cue_threshold_decay, target_cue_threshold_minimum=target_cue_threshold_minimum)
+	path_length = len(search_path)
+	final_step = search_path[-1]
+	if not search_successful:
+		# if no target is found with confidence passing confidence_high, select the target with the highest confidence during search and compare its confidence with confidence_low
+		max_logit = 0
+		final_step = None
+		path_length = 0
+		for i, search_step in enumerate(search_path):
+			if 'temp_detection_result' in search_step:
+				if search_step['temp_detection_result'][0] > max_logit:
+					max_logit = search_step['temp_detection_result'][0]
+					final_step = search_step
+					path_length = i+1
+		final_step['detection_result'] = final_step['temp_detection_result'][1]
+		if max_logit >= confidence_low:
+			search_successful = True
+	if visualize:
+		vis_path_length = path_length if search_successful else len(search_path)
+		visualize_search_path(image, search_path, vis_path_length, target_bbox, target_object_name, save_path)
+	del queue
+	return final_step, path_length, search_successful, all_valid_boxes
+
+
+
+def main(args):
+	args = parse_args(args)
+	vsm = VSM(args)
+
+	benchmark_folder = args.benchmark_folder
+
+	acc_list = []
+	search_path_length_list = []
+
+	for test_type in ['direct_attributes', 'relative_position']:
+		folder = os.path.join(benchmark_folder, test_type)
+		output_folder = None
+		if args.visualization:
+			output_folder =  os.path.join(args.output_path, test_type)
+			os.makedirs(output_folder, exist_ok=True)
+		image_files = filter(lambda file: '.json' not in file, os.listdir(folder))
+		for image_file in tqdm.tqdm(image_files):
+			image_path = os.path.join(folder, image_file)
+			annotation_path = image_path.split('.')[0] + '.json'
+			annotation = json.load(open(annotation_path))
+			bboxs = annotation['bbox']
+			object_names = annotation['target_object']
+
+			for i, (gt_bbox, object_name) in enumerate(zip(bboxs, object_names)):
+				image = Image.open(image_path).convert('RGB')
+				smallest_size = max(int(np.ceil(min(image.width, image.height)/args.minimum_size_scale)), args.minimum_size)
+				if args.visualization:
+					vis_path = os.path.join(output_folder, "{}_{}".format(image_file.split('.')[0],i))
+				else:
+					vis_path = None
+				final_step, path_length, search_successful, all_valid_boxes = visual_search(vsm, image, object_name, target_bbox=gt_bbox, smallest_size=smallest_size, confidence_high=args.confidence_high, confidence_low=args.confidence_low, target_cue_threshold=args.target_cue_threshold, target_cue_threshold_decay=args.target_cue_threshold_decay, target_cue_threshold_minimum=args.target_cue_threshold_minimum, save_path=vis_path, visualize=args.visualization)
+				if search_successful:
+					search_bbox = final_step['detection_result']
+					search_final_patch = final_step['bbox']
+					search_bbox[0] += search_final_patch[0]
+					search_bbox[1] += search_final_patch[1]
+					iou_i = iou(search_bbox, gt_bbox).item()
+					det_acc = 1.0 if iou_i > 0.5 else 0.0
+					acc_list.append(det_acc)
+					search_path_length_list.append(path_length)
+				else:
+					acc_list.append(0)
+					search_path_length_list.append(0)
+	print('Avg search path length:', np.mean([search_path_length_list[i] for i in range(len(search_path_length_list)) if acc_list[i]]))
+	print('Top 1 Acc:', np.mean(acc_list))
+
+if __name__ == "__main__":
+	main(sys.argv[1:])

vstar_bench_eval.py

@@ -0,0 +1,294 @@
+import argparse
+import os
+import json
+from tqdm import tqdm
+from collections import defaultdict
+from copy import deepcopy
+
+from PIL import Image
+import numpy as np
+import torch
+from torch.nn import CrossEntropyLoss
+
+from LLaVA.llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from LLaVA.llava.conversation import conv_templates, SeparatorStyle
+from LLaVA.llava.model.builder import load_pretrained_model
+from LLaVA.llava.utils import disable_torch_init
+from LLaVA.llava.mm_utils import get_model_name_from_path, KeywordsStoppingCriteria, tokenizer_image_object_token
+
+from visual_search import parse_args, VSM, visual_search
+
+def normalize_bbox(bbox, image_width, image_height):
+	normalized_bbox = [bbox[0]/image_width, bbox[1]/image_height, (bbox[0]+bbox[2])/image_width, (bbox[1]+bbox[3])/image_height]
+	normalized_bbox = [np.clip(_, 0, 1) for _ in normalized_bbox]
+	return normalized_bbox
+def expand2square(pil_img, background_color):
+	width, height = pil_img.size
+	if width == height:
+		return pil_img, 0, 0
+	elif width > height:
+		result = Image.new(pil_img.mode, (width, width), background_color)
+		result.paste(pil_img, (0, (width - height) // 2))
+		return result, 0, (width - height) // 2
+	else:
+		result = Image.new(pil_img.mode, (height, height), background_color)
+		result.paste(pil_img, ((height - width) // 2, 0))
+		return result, (height - width) // 2, 0
+
+class VQA_LLM:
+	def __init__(self, args):
+		disable_torch_init()
+		model_path = args.vqa_model_path
+		model_name = get_model_name_from_path(model_path)
+		model_name += 'llava'
+		model_base = None
+		device_map = "auto"
+		self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(model_path, model_base, model_name)
+		self.conv_type = args.conv_type
+
+	def get_patch(self, bbox, image_width, image_height, patch_size=224, patch_scale=None):
+		object_width = int(np.ceil(bbox[2]))
+		object_height = int(np.ceil(bbox[3]))
+
+		object_center_x = int(bbox[0] + bbox[2]/2)
+		object_center_y = int(bbox[1] + bbox[3]/2)
+
+		if patch_scale is None:
+			patch_width = max(object_width, patch_size)
+			patch_height = max(object_height, patch_size)
+		else:
+			patch_width = int(object_width*patch_scale)
+			patch_height = int(object_height*patch_scale)
+
+		left = max(0, object_center_x-patch_width//2)
+		right = min(left+patch_width, image_width)
+
+		top = max(0, object_center_y-patch_height//2)
+		bottom = min(top+patch_height, image_height)
+
+		return [left, top, right, bottom]
+	
+	def get_object_crop(self, image, bbox, patch_scale):
+		resized_bbox = self.get_patch(bbox, image.width, image.height, patch_scale=patch_scale)
+		object_crop = image.crop((resized_bbox[0], resized_bbox[1], resized_bbox[2], resized_bbox[3]))
+		object_crop = object_crop.resize((self.image_processor.crop_size['width'],self.image_processor.crop_size['height']))
+		object_crop = self.image_processor.preprocess(object_crop, return_tensors='pt')['pixel_values'][0]
+		return object_crop
+
+	@torch.inference_mode()
+	def free_form_inference(self, image, question, temperature=0, top_p=None, num_beams=1, max_new_tokens=200, object_crops=None, images_long=None, objects_long=None):
+		conv = conv_templates[self.conv_type].copy()
+		qs = DEFAULT_IMAGE_TOKEN + '\n' + question	
+		conv.append_message(conv.roles[0], qs)
+		conv.append_message(conv.roles[1], None)
+		prompt = conv.get_prompt()
+		stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+		keywords = [stop_str]
+		input_ids = tokenizer_image_object_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+		image_tensor = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+		stopping_criteria = KeywordsStoppingCriteria(keywords, self.tokenizer, input_ids)
+
+		output_ids = self.model.generate(
+			input_ids,
+			images=image_tensor.unsqueeze(0).half().cuda(),
+			object_features=object_crops.half().cuda() if object_crops is not None else None,
+			images_long = images_long,
+			objects_long = objects_long,
+			do_sample= True if temperature > 0 else False,
+			num_beams=num_beams,
+			temperature=temperature,
+			top_p = top_p,
+			max_new_tokens=max_new_tokens,
+			use_cache=True,
+			stopping_criteria=[stopping_criteria])
+			
+		input_token_len = input_ids.shape[1]
+		n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
+		if n_diff_input_output > 0:
+			print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+		outputs = self.tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
+		outputs = outputs.strip()
+		if outputs.endswith(stop_str):
+			outputs = outputs[:-len(stop_str)]
+		outputs = outputs.strip()
+		return outputs
+
+	@torch.inference_mode()
+	def multiple_choices_inference(self, image, question, options, object_crops=None, images_long=None, objects_long=None):
+		conv = conv_templates[self.conv_type].copy()
+		qs = DEFAULT_IMAGE_TOKEN + '\n' + question	
+		conv.append_message(conv.roles[0], qs)
+		conv.append_message(conv.roles[1], None)
+		prompt = conv.get_prompt()
+
+		question_input_ids = tokenizer_image_object_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+		image_tensor = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+
+		output_question = self.model(
+			question_input_ids,
+			use_cache=True,
+			images=image_tensor.unsqueeze(0).half().cuda(),
+			object_features=object_crops.half().cuda() if object_crops is not None else None,
+			images_long = images_long,
+			objects_long = objects_long)
+
+		question_logits = output_question.logits
+		question_past_key_values = output_question.past_key_values
+
+		loss_list = []
+
+		for option in options:
+			conv = conv_templates[self.conv_type].copy()
+			conv.append_message(conv.roles[0], qs)
+			conv.append_message(conv.roles[1], option)
+			full_prompt = conv.get_prompt()
+
+			full_input_ids = tokenizer_image_object_token(full_prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+			option_answer_input_ids = full_input_ids[:, question_input_ids.shape[1]:]
+
+			output_option = self.model(input_ids=option_answer_input_ids,
+								use_cache=True,
+								attention_mask=torch.ones(1, question_logits.shape[1]+option_answer_input_ids.shape[1], device=full_input_ids.device),
+								past_key_values=question_past_key_values)
+			
+			logits = torch.cat([question_logits[:, -1:], output_option.logits[:, :-1]], 1)
+
+			loss_fct = CrossEntropyLoss()
+			logits = logits.view(-1, self.model.config.vocab_size)
+			labels = option_answer_input_ids.view(-1)
+			loss = loss_fct(logits, labels)
+
+			loss_list.append(loss)
+
+		option_chosen = torch.stack(loss_list).argmin()
+
+		return option_chosen.cpu().item()
+
+
+def eval_model(args):
+	# init VQA LLM
+	vqa_llm = VQA_LLM(args)
+	# init VSM
+	vsm_args = parse_args({})
+	vsm_args.version = args.vsm_model_path
+	vsm = VSM(vsm_args)
+
+	results = {}
+	per_type_acc = defaultdict(list)
+	all_acc = []
+
+	missing_objects_msg = "Sorry, I can not answer the question. Some visual information about the following objects is missing or unclear:"
+	focus_msg = "Additional visual information to focus on: "
+	for test_type in ['direct_attributes', 'relative_position']:
+		results[test_type] = []
+		folder = os.path.join(args.benchmark_folder, test_type)
+		image_files = list(filter(lambda file: '.json' not in file, os.listdir(folder)))
+		for image_file in tqdm(image_files):
+			result_single_sample = {}
+			image_path = os.path.join(folder, image_file)
+			annotation_path = image_path.split('.')[0] + '.json'
+			image = Image.open(image_path).convert('RGB')
+			annotation = json.load(open(annotation_path))
+			image, _, _ = expand2square(image, tuple(int(x*255) for x in vqa_llm.image_processor.image_mean))
+			
+			question = annotation['question']
+			# generate free-form response to check whether visual search needs to be activated
+			prediction = vqa_llm.free_form_inference(image, question)
+			missing_objects = []
+			if missing_objects_msg in prediction:
+				missing_objects = prediction.split(missing_objects_msg)[-1]
+				if missing_objects.endswith('.'):
+					missing_objects = missing_objects[:-1]
+				missing_objects = missing_objects.split(',')
+				missing_objects = [missing_object.strip() for missing_object in missing_objects]
+
+			search_result = []
+			if len(missing_objects) > 0:
+				# visual search
+				for object_name in missing_objects:
+					image = Image.open(image_path).convert('RGB')
+					smallest_size = max(int(np.ceil(min(image.width, image.height)/args.minimum_size_scale)), args.minimum_size)
+					final_step, path_length, search_successful, all_valid_boxes = visual_search(vsm, image, object_name, target_bbox=None, smallest_size=smallest_size)
+					if all_valid_boxes is not None:
+						# might exist multiple target instances
+						for search_bbox in all_valid_boxes:
+							search_final_patch = final_step['bbox']
+							search_bbox[0] += search_final_patch[0]
+							search_bbox[1] += search_final_patch[1]
+							search_result.append({'bbox':search_bbox.tolist(),'name':object_name})
+					else:
+						search_bbox = final_step['detection_result']
+						search_final_patch = final_step['bbox']
+						search_bbox[0] += search_final_patch[0]
+						search_bbox[1] += search_final_patch[1]
+						search_result.append({'bbox':search_bbox.tolist(),'name':object_name})
+			# predict the multiple-choice option
+			options = annotation['options']
+			image = Image.open(image_path).convert('RGB')
+			if len(missing_objects) > 0:
+				object_names = [_['name'] for _ in search_result]
+				bboxs = deepcopy([_['bbox'] for _ in search_result])
+				if len(object_names) <= 2:
+					images_long = [False]
+					objects_long = [True]*len(object_names)
+				else:
+					images_long = [False]
+					objects_long = [False]*len(object_names)
+				object_crops = []
+				for bbox in bboxs:
+					object_crop = vqa_llm.get_object_crop(image, bbox, patch_scale=1.2)
+					object_crops.append(object_crop)
+				object_crops = torch.stack(object_crops, 0)
+				image, left, top = expand2square(image, tuple(int(x*255) for x in vqa_llm.image_processor.image_mean))
+				bbox_list = []
+				for bbox in bboxs:
+					bbox[0] += left
+					bbox[1] += top
+					bbox_list.append(bbox)
+				bbox_list = [normalize_bbox(bbox, image.width, image.height) for bbox in bbox_list]
+				cur_focus_msg = focus_msg
+				for i, (object_name, bbox) in enumerate(zip(object_names, bbox_list)):
+					cur_focus_msg = cur_focus_msg + "{} <object> at location [{:.3f},{:.3f},{:.3f},{:.3f}]".format(object_name, bbox[0], bbox[1], bbox[2], bbox[3])
+					if i != len(bbox_list)-1:
+						cur_focus_msg = cur_focus_msg+"; "
+					else:
+						cur_focus_msg = cur_focus_msg +'.'
+				question_with_focus = cur_focus_msg+"\n"+question
+				option_chosen = vqa_llm.multiple_choices_inference(image, question_with_focus, options, object_crops, images_long=images_long, objects_long=objects_long)
+			else:
+				option_chosen = vqa_llm.multiple_choices_inference(image, question, options)
+
+			correct = 1 if option_chosen==0 else 0
+			per_type_acc[test_type].append(correct)
+			all_acc.append(correct)
+
+			result_single_sample['question'] = question
+			result_single_sample['options'] = options
+			result_single_sample['image'] = image_file
+			result_single_sample['prediction_freeform'] = prediction
+			result_single_sample['missing_objects'] = missing_objects
+			result_single_sample['search_result'] = search_result	
+			result_single_sample['option_chosen'] = option_chosen
+			result_single_sample['correct'] = correct
+			results[test_type].append(result_single_sample)
+
+		print(test_type, np.mean(per_type_acc[test_type]))
+
+	print(np.mean(all_acc))
+
+	with open(args.output_path, 'w') as f:
+		json.dump(results, f, indent=4)
+
+if __name__ == "__main__":
+	parser = argparse.ArgumentParser()
+	parser.add_argument("--vqa-model-path", type=str, default="craigwu/seal_vqa_7b")
+	parser.add_argument("--vqa-model-base", type=str, default=None)
+	parser.add_argument("--conv_type", default="v1", type=str,)
+	parser.add_argument("--benchmark-folder", type=str, default="vstar_bench")
+	parser.add_argument("--vsm-model-path", type=str, default="craigwu/seal_vsm_7b")
+	parser.add_argument("--output-path", type=str, default="eval_result.json")
+	parser.add_argument("--minimum_size_scale", default=4.0, type=float, help="minimum sub-image scale for the termination of search")
+	parser.add_argument("--minimum_size", default=224, type=int, help="minimum sub-image size for the termination of search")
+
+	args = parser.parse_args()
+	eval_model(args)