""" TStarSearcher: Comprehensive Video Frame Search Tool This script allows searching for specific objects within a video using YOLO object detection and GPT-4 for question-answering. It leverages the TStar framework's universal Grounder, YOLO interface, and video searcher to identify relevant frames and answer questions based on the detected objects. Usage: python tstar_searcher.py --video_path path/to/video.mp4 --question "Your question here" --options "A) Option1\nB) Option2\nC) Option3\nD) Option4" """ import os import sys import cv2 import torch import copy import logging import argparse import numpy as np import matplotlib.pyplot as plt from dataclasses import dataclass, field from typing import List, Optional, Tuple from decord import VideoReader, cpu from scipy.interpolate import UnivariateSpline # Import custom TStar interfaces from TStar.interface_llm import TStarUniversalGrounder from TStar.interface_owl import OWLInterface from TStar.interface_searcher import TStarSearcher from TStar.utils import save_as_gif # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', handlers=[ logging.StreamHandler(sys.stdout) ] ) logger = logging.getLogger(__name__) class TStarFramework: """ Main class for performing object-based frame search and question-answering in a video. """ def __init__( self, video_path: str, heuristic_scorer: OWLInterface, grounder: TStarUniversalGrounder, question: str, options: str, search_nframes: int = 8, grid_rows: int = 4, grid_cols: int = 4, output_dir: str = './output', confidence_threshold: float = 0.6, search_budget: int = 1000, prefix: str = 'stitched_image', config_path: Optional[str] = None, checkpoint_path: Optional[str] = None, device: str = "cuda:0" ): """ Initialize VideoSearcher. Args: video_path (str): Path to the input video file. yolo_scorer (YoloV5Interface): YOLO interface instance. grounder (TStarUniversalGrounder): Universal Grounder instance. question (str): The question for question-answering. options (str): Multiple-choice options for the question. search_nframes (int, optional): Number of top frames to return. Default is 8. grid_rows (int, optional): Number of rows in the image grid. Default is 4. grid_cols (int, optional): Number of columns in the image grid. Default is 4. output_dir (str, optional): Directory to save outputs. Default is './output'. confidence_threshold (float, optional): YOLO detection confidence threshold. Default is 0.6. search_budget (int, optional): Maximum number of frames to process during search. Default is 1000. prefix (str, optional): Prefix for output filenames. Default is 'stitched_image'. config_path (str, optional): Path to the YOLO configuration file. Default is None. checkpoint_path (str, optional): Path to the YOLO model checkpoint. Default is None. device (str, optional): Device for model inference (e.g., "cuda:0" or "cpu"). Default is "cuda:0". """ self.video_path = video_path self.yolo_scorer = heuristic_scorer self.grounder = grounder self.question = question self.options = options self.search_nframes = search_nframes self.grid_rows = grid_rows self.grid_cols = grid_cols self.output_dir = output_dir self.confidence_threshold = confidence_threshold self.search_budget = search_budget self.prefix = prefix self.config_path = config_path self.checkpoint_path = checkpoint_path self.device = device # Ensure the output directory exists os.makedirs(self.output_dir, exist_ok=True) logger.info("VideoSearcher initialized successfully.") self.results = {} def run(self): """ Execute the complete video search and question-answering process. """ # Use Grounder to get target and cue objects target_objects, cue_objects = self.get_grounded_objects() # Initialize TStarSearcher video_searcher = TStarSearcher( video_path=self.video_path, target_objects=target_objects, cue_objects=cue_objects, search_nframes=self.search_nframes, image_grid_shape=(self.grid_rows, self.grid_cols), output_dir=self.output_dir, confidence_threshold=self.confidence_threshold, search_budget=self.search_budget, prefix=self.prefix, owl=self.yolo_scorer ) logger.info(f"TStarSearcher initialized successfully for video {self.video_path}.") # Perform search all_frames, time_stamps = self.perform_search(video_searcher) # Save retrieved frames self.save_frames(all_frames, time_stamps) self.save_searching_iters(video_searcher) # Plot and save score distribution self.plot_and_save_scores(video_searcher) # Perform question-answering on retrieved frames answer = self.perform_qa(all_frames) print("QA Answer:", answer) logger.info("VideoSearcher completed successfully.") def get_grounded_objects(self) -> Tuple[List[str], List[str]]: """ Use Grounder to obtain target and cue objects. Returns: Tuple[List[str], List[str]]: Lists of target objects and cue objects. """ # Example code; should be implemented based on Grounder's interface # For example: target_objects, cue_objects = self.grounder.inference_query_grounding( video_path=self.video_path, question=self.question ) # Here, assuming fixed target and cue objects # target_objects = ["couch"] # Target objects to find # cue_objects = ["TV", "chair"] # Cue objects logger.info(f"Target objects: {target_objects}") logger.info(f"Cue objects: {cue_objects}") self.results["Searching_Objects"] = {"target_objects": target_objects, "cue_objects": cue_objects} return target_objects, cue_objects def perform_search(self, video_searcher: TStarSearcher) -> Tuple[List[np.ndarray], List[float]]: """ Execute the frame search process and retrieve relevant frames and timestamps. Args: video_searcher (TStarSearcher): Instance of TStarSearcher. Returns: Tuple[List[np.ndarray], List[float]]: List of frames and their corresponding timestamps. """ all_frames, time_stamps = video_searcher.search_with_visualization() logger.info(f"Found {len(all_frames)} frames, timestamps: {time_stamps}") self.results['timestamps'] = time_stamps return all_frames, time_stamps def perform_qa(self, frames: List[np.ndarray]) -> str: """ Perform question-answering on the retrieved frames. Args: frames (List[np.ndarray]): List of frames to analyze. Returns: str: Answer generated by VLM. """ answer = self.grounder.inference_qa( frames=frames, question=self.question, options=self.options ) self.results['answer'] = answer return answer def plot_and_save_scores(self, video_searcher: TStarSearcher): """ Plot the score distribution and save the plot. Args: video_searcher (TStarSearcher): Instance of TStarSearcher. """ plot_path = os.path.join(self.output_dir, "score_distribution.png") video_searcher.plot_score_distribution(save_path=plot_path) logger.info(f"Score distribution plot saved to {plot_path}") def save_frames(self, frames: List[np.ndarray], timestamps: List[float]): """ Save the retrieved frames as image files. Args: frames (List[np.ndarray]): List of frames to save. timestamps (List[float]): Corresponding timestamps of the frames. """ for idx, (frame, timestamp) in enumerate(zip(frames, timestamps)): frame_path = os.path.join( self.output_dir, f"frame_{idx}_at_{timestamp:.2f}s.jpg" ) cv2.imwrite(frame_path, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) logger.info(f"Saved frame to {frame_path}") def save_searching_iters(self, video_searcher, video_ids=[]): # # 定义 resize 操作，目标大小为 (640, 640) # resize_transform = T.Resize((1024, 1024)) # resized_frames_tensor = resize_transform(resized_frames_tensor) image_grid_iters = video_searcher.image_grid_iters # iters, b, image # b = 1 for v1 detect_annotot_iters = video_searcher.detect_annotot_iters # iters, b, image detect_bbox_iters = video_searcher.detect_bbox_iters #iters, b, n_objects, xxyy, fps = 1 # 设置帧率为 2 for b in range(len(image_grid_iters[0])): images = [image_grid_iter[b] for image_grid_iter in image_grid_iters] anno_images = [detect_annotot_iter[b] for detect_annotot_iter in detect_annotot_iters] frame_size = (anno_images[0].shape[1], anno_images[0].shape[0]) # 获取图像大小 (宽度, 高度) # 设置视频的参数 video_id=self.video_path.split("/")[-1].split(".")[0] output_video_path = os.path.join(self.output_dir, f"{video_id}.gif") # 视频保存路径 save_as_gif(images=anno_images, output_gif_path=output_video_path) def initialize_TStar_Scorer( heuristic: str, device: str ) -> OWLInterface: """ Initialize the YOLO object detection model. Args: config_path (str): Path to the YOLO configuration file. checkpoint_path (str): Path to the YOLO model checkpoint. device (str): Device for model inference (e.g., "cuda:0"). Returns: YoloWorldInterface: Initialized YOLO interface instance. Raises: FileNotFoundError: If the configuration file or checkpoint file is not found. """ model_choice = 'owl_model' if model_choice == 'owl_model': model_name="google/owlvit-base-patch32" owl_interface = OWLInterface( config_path = model_name, checkpoint_path=None, device="cuda:0" ) logger.info("YoloWorldInterface initialized successfully.") return owl_interface def run_tstar( video_path: str, question: str, options: str, grounder: str, heuristic: str, openai_api_key: str, device: str = "cuda:0", search_nframes: int = 8, grid_rows: int = 4, grid_cols: int = 4, confidence_threshold: float = 0.6, search_budget: float = 0.5, output_dir: str = './output', ): """ Executes the TStar video frame search and QA process. Args: video_path (str): Path to the input video file. question (str): Question for video content QA. options (str): Multiple-choice options for the question. config_path (str): Path to the YOLO configuration file. checkpoint_path (str): Path to the YOLO model checkpoint. device (str): Device for model inference (e.g., "cuda:0" or "cpu"). search_nframes (int): Number of top frames to return. grid_rows (int): Number of rows in the image grid. grid_cols (int): Number of columns in the image grid. confidence_threshold (float): YOLO detection confidence threshold. search_budget (float): Maximum ratio of frames to process during search. output_dir (str): Directory to save outputs. prefix (str): Prefix for output filenames. Returns: dict: Results containing detected objects, timestamps, and the QA answer. """ # Initialize Grounder and YOLO grounder = TStarUniversalGrounder(backend="gpt4", model_name="gpt-4o", gpt4_api_key=openai_api_key) TStar_Scorer = initialize_TStar_Scorer( heuristic=heuristic, device=device ) # Initialize and run the search framework searcher = TStarFramework( grounder=grounder, heuristic_scorer=TStar_Scorer, video_path=video_path, question=question, options=options, search_nframes=search_nframes, grid_rows=grid_rows, grid_cols=grid_cols, output_dir=output_dir, confidence_threshold=confidence_threshold, search_budget=search_budget, device=device ) searcher.run() return { "Grounding Objects": searcher.results.get('Searching_Objects', []), "Frame Timestamps": searcher.results.get('timestamps', []), "Answer": searcher.results.get('answer', "No answer generated") }