from typing import Optional, Tuple, Union, List import numpy as np import PIL from PIL.Image import Image import supervision as sv import torch from torch import nn from transformers import OwlViTProcessor, OwlViTForObjectDetection, OwlViTVisionModel from transformers.models.owlvit.modeling_owlvit import center_to_corners_format, box_iou, generalized_box_iou, OwlViTObjectDetectionOutput from sam_extension.pipeline.base import Pipeline, Output class OwlViTVisionEncoderPipeline(Pipeline): def __init__(self, vision_model, layer_norm, processor, device='cuda', *args, **kwargs): super().__init__(*args, **kwargs) self.vision_model = vision_model self.layer_norm = layer_norm self.processor = processor self.device = device torch.cuda.empty_cache() @classmethod def from_pretrained(cls, model_type, device='cuda', *args, **kwargs): owlvit_for_object_detection = OwlViTForObjectDetection.from_pretrained(model_type).to(device) processor = OwlViTProcessor.from_pretrained(model_type) return cls(owlvit_for_object_detection.owlvit.vision_model, owlvit_for_object_detection.layer_norm, processor, device, *args, **kwargs) def process_image(self, image:Image): image = self.processor(images=image, return_tensors="pt").pixel_values.to(self.device) return image @torch.no_grad() def forward( self, pixel_values: Union[torch.FloatTensor, Image] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> torch.FloatTensor: if isinstance(pixel_values, Image): pixel_values = self.process_image(pixel_values) pixel_values = pixel_values.to(self.device) vision_outputs = self.vision_model( pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) # Get image embeddings last_hidden_state = vision_outputs[0] image_embeds = self.vision_model.post_layernorm(last_hidden_state) new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0))) class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size) # Merge image embedding with class tokens image_embeds = image_embeds[:, 1:, :] * class_token_out image_embeds = self.layer_norm(image_embeds) # Resize to [batch_size, num_patches, num_patches, hidden_size] new_size = ( image_embeds.shape[0], int(np.sqrt(image_embeds.shape[1])), int(np.sqrt(image_embeds.shape[1])), image_embeds.shape[-1], ) image_embeds = image_embeds.reshape(new_size) return image_embeds class OwlViTDecoderPipeline(Pipeline): prompt_template: str = 'a photo of a ' def __init__(self, owlvit_text, text_projection, class_head, box_head, processor, device='cuda', *args, **kwargs): super().__init__(*args, **kwargs) self.owlvit_text = owlvit_text self.text_projection = text_projection self.class_head = class_head self.box_head = box_head self.sigmoid = nn.Sigmoid() self.processor = processor self.device = device torch.cuda.empty_cache() @classmethod def from_pretrained(cls, model_type, device='cuda', *args, **kwargs): owlvit_for_object_detection = OwlViTForObjectDetection.from_pretrained(model_type).to(device) processor = OwlViTProcessor.from_pretrained(model_type) return cls(owlvit_for_object_detection.owlvit.text_model, owlvit_for_object_detection.owlvit.text_projection, owlvit_for_object_detection.class_head, owlvit_for_object_detection.box_head, processor, device, *args, **kwargs) def set_template(self, template: str): self.prompt_template = template def process_text(self, text:List, use_template:bool = True): if use_template: text = [[self.prompt_template+i for i in text[0]]] inputs = self.processor(text=text, return_tensors="pt") return inputs def normalize_grid_corner_coordinates(self, feature_map: torch.FloatTensor): # Computes normalized xy corner coordinates from feature_map. if not feature_map.ndim == 4: raise ValueError("Expected input shape is [batch_size, num_patches, num_patches, hidden_dim]") device = feature_map.device num_patches = feature_map.shape[1] box_coordinates = np.stack( np.meshgrid(np.arange(1, num_patches + 1), np.arange(1, num_patches + 1)), axis=-1 ).astype(np.float32) box_coordinates /= np.array([num_patches, num_patches], np.float32) # Flatten (h, w, 2) -> (h*w, 2) box_coordinates = box_coordinates.reshape( box_coordinates.shape[0] * box_coordinates.shape[1], box_coordinates.shape[2] ) box_coordinates = torch.from_numpy(box_coordinates).to(device) return box_coordinates def compute_box_bias(self, feature_map: torch.FloatTensor) -> torch.FloatTensor: # The box center is biased to its position on the feature grid box_coordinates = self.normalize_grid_corner_coordinates(feature_map) box_coordinates = torch.clip(box_coordinates, 0.0, 1.0) # Unnormalize xy box_coord_bias = torch.log(box_coordinates + 1e-4) - torch.log1p(-box_coordinates + 1e-4) # The box size is biased to the patch size box_size = torch.full_like(box_coord_bias, 1.0 / feature_map.shape[-2]) box_size_bias = torch.log(box_size + 1e-4) - torch.log1p(-box_size + 1e-4) # Compute box bias box_bias = torch.cat([box_coord_bias, box_size_bias], dim=-1) return box_bias def box_predictor( self, image_feats: torch.FloatTensor, feature_map: torch.FloatTensor, ) -> torch.FloatTensor: """ Args: image_feats: Features extracted from the image, returned by the `image_text_embedder` method. feature_map: A spatial re-arrangement of image_features, also returned by the `image_text_embedder` method. Returns: pred_boxes: List of predicted boxes (cxcywh normalized to 0, 1) nested within a dictionary. """ # Bounding box detection head [batch_size, num_boxes, 4]. pred_boxes = self.box_head(image_feats) # Compute the location of each token on the grid and use it to compute a bias for the bbox prediction pred_boxes += self.compute_box_bias(feature_map) pred_boxes = self.sigmoid(pred_boxes) return pred_boxes def class_predictor( self, image_feats: torch.FloatTensor, query_embeds: Optional[torch.FloatTensor] = None, query_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.FloatTensor]: """ Args: image_feats: Features extracted from the `image_text_embedder`. query_embeds: Text query embeddings. query_mask: Must be provided with query_embeddings. A mask indicating which query embeddings are valid. """ (pred_logits, image_class_embeds) = self.class_head(image_feats, query_embeds, query_mask) return (pred_logits, image_class_embeds) def image_text_embedder( self, input_ids: torch.Tensor, image_embeds: torch.FloatTensor, attention_mask: torch.Tensor, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, ) -> Tuple[torch.FloatTensor]: # Encode text and image text_outputs = self.owlvit_text( input_ids=input_ids, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, ) text_embeds = text_outputs[1] text_embeds = self.text_projection(text_embeds) text_embeds = text_embeds / torch.linalg.norm(text_embeds, ord=2, dim=-1, keepdim=True) return (text_embeds, image_embeds, text_outputs) def embed_image_query( self, query_image_features: torch.FloatTensor, query_feature_map: torch.FloatTensor ) -> torch.FloatTensor: _, class_embeds = self.class_predictor(query_image_features) pred_boxes = self.box_predictor(query_image_features, query_feature_map) pred_boxes_as_corners = center_to_corners_format(pred_boxes) # Loop over query images best_class_embeds = [] best_box_indices = [] pred_boxes_device = pred_boxes_as_corners.device for i in range(query_image_features.shape[0]): each_query_box = torch.tensor([[0, 0, 1, 1]], device=pred_boxes_device) each_query_pred_boxes = pred_boxes_as_corners[i] ious, _ = box_iou(each_query_box, each_query_pred_boxes) # If there are no overlapping boxes, fall back to generalized IoU if torch.all(ious[0] == 0.0): ious = generalized_box_iou(each_query_box, each_query_pred_boxes) # Use an adaptive threshold to include all boxes within 80% of the best IoU iou_threshold = torch.max(ious) * 0.8 selected_inds = (ious[0] >= iou_threshold).nonzero() if selected_inds.numel(): selected_embeddings = class_embeds[i][selected_inds[0]] mean_embeds = torch.mean(class_embeds[i], axis=0) mean_sim = torch.einsum("d,id->i", mean_embeds, selected_embeddings) best_box_ind = selected_inds[torch.argmin(mean_sim)] best_class_embeds.append(class_embeds[i][best_box_ind]) best_box_indices.append(best_box_ind) if best_class_embeds: query_embeds = torch.stack(best_class_embeds) box_indices = torch.stack(best_box_indices) else: query_embeds, box_indices = None, None return query_embeds, box_indices, pred_boxes @torch.no_grad() def forward( self, image_embeds: torch.FloatTensor, input_ids: Optional[torch.Tensor] = None, text: Optional[List] = None, attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> OwlViTObjectDetectionOutput: if text is not None: inputs = self.process_text(text) input_ids = inputs.input_ids.to(self.device) attention_mask = inputs.attention_mask.to(self.device) input_ids = input_ids.to(self.device) image_embeds = image_embeds.to(self.device) attention_mask = attention_mask.to(self.device) output_attentions = output_attentions if output_attentions is not None else False output_hidden_states = ( output_hidden_states if output_hidden_states is not None else False ) return_dict = return_dict if return_dict is not None else True # Embed images and text queries query_embeds, feature_map, text_outputs = self.image_text_embedder( input_ids=input_ids, image_embeds=image_embeds, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, ) # Text and vision model outputs batch_size, num_patches, num_patches, hidden_dim = feature_map.shape image_feats = torch.reshape(feature_map, (batch_size, num_patches * num_patches, hidden_dim)) # Reshape from [batch_size * max_text_queries, hidden_dim] -> [batch_size, max_text_queries, hidden_dim] max_text_queries = input_ids.shape[0] // batch_size query_embeds = query_embeds.reshape(batch_size, max_text_queries, query_embeds.shape[-1]) # If first token is 0, then this is a padded query [batch_size, num_queries]. input_ids = input_ids.reshape(batch_size, max_text_queries, input_ids.shape[-1]) query_mask = input_ids[..., 0] > 0 # Predict object classes [batch_size, num_patches, num_queries+1] (pred_logits, class_embeds) = self.class_predictor(image_feats, query_embeds, query_mask) # Predict object boxes pred_boxes = self.box_predictor(image_feats, feature_map) if not return_dict: output = ( pred_logits, pred_boxes, query_embeds, feature_map, class_embeds, text_outputs.to_tuple(), None, ) output = tuple(x for x in output if x is not None) return output return OwlViTObjectDetectionOutput( image_embeds=feature_map, text_embeds=query_embeds, pred_boxes=pred_boxes.cpu(), logits=pred_logits.cpu(), class_embeds=class_embeds, text_model_output=text_outputs, vision_model_output=None, ) def owlvit_visualize(self, image: Image, texts: List, owlvit_objectdetection_output: OwlViTObjectDetectionOutput, score_threshold: float = 0.1, pil=True): target_sizes = torch.Tensor([image.size[::-1]]) # Convert outputs (bounding boxes and class logits) to COCO API results = self.processor.post_process(outputs=owlvit_objectdetection_output, target_sizes=target_sizes) text = texts[0] boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"] boxes_np = [] labels_list = [] # Print detected objects and rescaled box coordinates for box, score, label in zip(boxes, scores, labels): box = [int(i) for i in box.tolist()] if score >= score_threshold: labels_list.append(f"{text[label]} {round(score.item(), 3)}") boxes_np.append(box) print(f"Detected {text[label]} with confidence {round(score.item(), 3)} at location {box}") boxes_np = np.array(boxes_np) detections = sv.Detections(xyxy=boxes_np) image_np = np.uint8(image)[:, :, ::-1] box_annotator = sv.BoxAnnotator() annotated_frame = box_annotator.annotate(scene=image_np.copy(), detections=detections, labels=labels_list) if pil: return PIL.Image.fromarray(annotated_frame[:, :, ::-1]) else: return annotated_frame[:, :, ::-1]