import io from typing import Any, Dict, List import cv2 import tempfile import numpy as np import requests import torch from PIL import Image from transformers import AutoTokenizer, XCLIPModel, XCLIPProcessor from huggingface_hub import logging from concurrent.futures import ThreadPoolExecutor, as_completed import timeit device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class EndpointHandler: def __init__(self, path=""): # Preload all the elements you are going to need at inference. # pseudo # self.model = load_model(path) model_id = "microsoft/xclip-base-patch16-zero-shot" # self.device = "cuda" if torch.cuda.is_available() else "cpu" self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.processor = XCLIPProcessor.from_pretrained(path) self.model = XCLIPModel.from_pretrained(path).to(self.device) self.tokenizer = AutoTokenizer.from_pretrained(path) logging.set_verbosity_debug() self.logger = logging.get_logger(__name__) # Check if CUDA (GPU support) is available if torch.cuda.is_available(): self.logger.info("GPU is available for inference.") self.logger.info(f"Using {torch.cuda.get_device_name(0)}") else: self.logger.info("GPU is not available, using CPU for inference.") def download_video_as_bytes(self, url: str) -> (bytes, dict): """ Download a video from a given URL, load it in RAM, and return it as bytes. Parameters: - url (str): The URL of the video to download. Returns: - bytes or None: The video content as bytes if successful, None otherwise. - dict or None: The video download headers if succesful, None otherwise. """ try: response = requests.get(url) response.raise_for_status() # Raise an exception for HTTP errors return response.content, response.headers except requests.RequestException as e: print(f"Error downloading the video: {e}") return None, None def extract_evenly_spaced_frames_from_bytes( self, video_bytes: bytes, num_frames: int = 32 ) -> list: # Write bytes to a temporary file with tempfile.NamedTemporaryFile(delete=True, suffix=".mp4") as temp_video: temp_video.write(video_bytes) temp_video.flush() # Create a VideoCapture object using the temporary file's name vidcap = cv2.VideoCapture(temp_video.name) # Get the total number of frames in the video total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) # Calculate the interval at which frames should be extracted interval = total_frames // num_frames frames = [] for i in range(num_frames): # Set the video position to the next frame to be captured vidcap.set(cv2.CAP_PROP_POS_FRAMES, i * interval) # Read the frame success, image = vidcap.read() # If successfully read, add the frame to the list if success: frames.append(image) return frames def preprocess_frames(self, video_frames): """ Define a preprocessing function to convert video frames into a format suitable for the model """ frames = np.array(video_frames) # Use the XCLIP Processor to preprocess the frames inputs = self.processor( text=None, videos=list(frames), return_tensors="pt", padding=True ).to(self.device) return inputs def embed_frames_with_xclip_processing(self, frames): # Initialize an empty list to store the frame embeddings self.logger.info("Preprocessing frames.") frame_preprocessed = self.preprocess_frames(frames) # Pass the preprocessed frame through the model to get the frame embeddings self.logger.info("Getting video features.") frame_embedding = self.model.get_video_features(**frame_preprocessed) # Stack the list of frame embeddings into a single tensor # self.logger.info("Stacking embeddings into a single tensor.") # tensor = torch.stack(frame_embedding) # detach text emb from graph, move to CPU, and convert to numpy array # self.logger.info("Squeezing tensor") # batch_emb = frame_embedding.squeeze(0) # Check the shape of the tensor self.logger.info(f"Shape of the batch_emb tensor: {frame_embedding.shape}") # Normalize the embeddings if it's a 2D tensor if frame_embedding.dim() == 2: self.logger.info("Normalizing embeddings") batch_emb = torch.nn.functional.normalize(frame_embedding, p=2, dim=1) else: self.logger.info("Skipping normalization due to tensor shape") batch_emb = frame_embedding.squeeze(0) self.logger.info("Converting into numpy array") batch_emb = batch_emb.cpu().detach().numpy() self.logger.info("Converting to list") batch_emb = batch_emb.tolist() self.logger.info("Returning batch_emb list") return batch_emb def process_video(self, video_url, video_metadata): try: self.logger.info("Downloading video as bytes.") video_bytes, video_headers = self.download_video_as_bytes(video_url) self.logger.info("Extracting frames.") frames = self.extract_evenly_spaced_frames_from_bytes( video_bytes, num_frames=32 ) self.logger.info("Embedding frames with Xclip.") frame_embeddings = self.embed_frames_with_xclip_processing(frames) video_metadata["url"] = video_url self.logger.info("Returning embeddings and metadata.") return frame_embeddings, video_metadata except Exception as e: print(e) return None, None, None def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: """ Process the input data based on its type and return the embeddings. This method accepts a dictionary with a 'process_type' key that can be either 'images' or 'text'. If 'process_type' is 'images', the method expects a list of image URLs under the 'images_urls' key. It downloads and processes these images, and returns their embeddings. If 'process_type' is 'text', the method expects a string query under the 'query' key. It processes this text and returns its embedding. Parameters: - data: Dict[str, Any] A dictionary containing the data to be processed. It must include a 'process_type' key with value either 'images' or 'text'. If 'process_type' is 'images', data should also include 'images_urls' key with a list of image URLs. If 'process_type' is 'text', data should also include 'query' key with a string query. Returns: - List[Dict[str, Any]] A list of dictionaries, each containing the embeddings of the processed data. If an error occurs during processing, the dictionary will include an 'error' key with the error message. Raises: - ValueError: If the 'process_type' key is not present in data, or if the required keys for 'images' or 'text' are not present or are of the wrong type. """ if data["process_type"] == "videos": try: if "videos_urls" not in data or not isinstance( data["videos_urls"], list ): raise ValueError( "Data must contain 'videos_urls' key with a list of videos urls." ) batch_size = 10 if "batch_size" in data: batch_size = int(data["batch_size"]) # Download and process the videos processed_video_embeddings = [] processed_videos_metadata = [] for i in range(0, len(data["videos_urls"]), batch_size): videos_urls = data["videos_urls"][i : i + batch_size] videos_metadata = data["videos_metadata"][i : i + batch_size] with ThreadPoolExecutor() as executor: futures = [ executor.submit(self.process_video, url, metadata) for url, metadata in zip(videos_urls, videos_metadata) ] for future in as_completed(futures): frame_embeddings, video_metadata = future.result() if frame_embeddings is not None: processed_video_embeddings.append(frame_embeddings) self.logger.info("Finished appending video embedding") processed_metadata = { "text": video_metadata["caption"], "source": video_metadata["url"], "source_type": "video_frames", **video_metadata, } processed_videos_metadata.append(processed_metadata) self.logger.info("Finished appending video metadata") self.logger.info(f"Finished processing batch {i}") # Return the embeddings self.logger.info("Returning embeddings and metadata of all batches") return { "embeddings": processed_video_embeddings, "metadata": processed_videos_metadata, } except Exception as e: print(f"Error during videos processing: {str(e)}") return {"embeddings": [], "error": str(e)} elif data["process_type"] == "text": if "query" not in data or not isinstance(data["query"], str): raise ValueError("Data must contain 'query' key which is a str.") query = data["query"] inputs = self.tokenizer(query, return_tensors="pt").to(self.device) text_emb = self.model.get_text_features(**inputs) # detach text emb from graph, move to CPU, and convert to numpy array text_emb = text_emb.detach().cpu().numpy() # calculate value to normalize each vector by and normalize them norm_factor = np.linalg.norm(text_emb, axis=1) text_emb = text_emb.T / norm_factor # transpose back to (21, 512) text_emb = text_emb.T # Converting tensor to list for JSON response text_emb_list = text_emb.tolist() return {"embeddings": text_emb_list} else: print( f"Error during CLIP endpoint processing: data['process_type']: {data['process_type']} neither 'images' or 'text'" ) return {"embeddings": [], "error": str(e)} # pseudo # self.model(input)