import gradio as gr from PIL import Image, ImageDraw import torch from transformers import OwlViTProcessor, OwlViTForObjectDetection, OwlViTModel, OwlViTImageProcessor from transformers.image_transforms import center_to_corners_format from transformers.models.owlvit.modeling_owlvit import box_iou from functools import partial import numpy as np # from utils import iou processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32") model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32") from transformers.models.owlvit.modeling_owlvit import OwlViTImageGuidedObjectDetectionOutput, OwlViTClassPredictionHead def classpredictionhead_box_forward( self, image_embeds, query_indice, query_mask, ): image_class_embeds = self.dense0(image_embeds) # Normalize image and text features image_class_embeds = image_class_embeds / (torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6) query_embeds = image_class_embeds[0, query_indice].unsqueeze(0).unsqueeze(0) # query_embeds = query_embeds / (torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6) # Get class predictions pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds) # Apply a learnable shift and scale to logits logit_shift = self.logit_shift(image_embeds) logit_scale = self.logit_scale(image_embeds) logit_scale = self.elu(logit_scale) + 1 pred_logits = (pred_logits + logit_shift) * logit_scale if query_mask is not None: if query_mask.ndim > 1: query_mask = torch.unsqueeze(query_mask, dim=-2) pred_logits = pred_logits.to(torch.float64) pred_logits = torch.where(query_mask == 0, -1e6, pred_logits) pred_logits = pred_logits.to(torch.float32) return (pred_logits, image_class_embeds) def class_predictor( self, image_feats, query_indice=None, query_mask=None, ): (pred_logits, image_class_embeds) = self.class_head.classpredictionhead_box_forward(image_feats, query_indice, query_mask) return (pred_logits, image_class_embeds) def get_max_iou_indice(target_pred_boxes, query_box, target_sizes): boxes = center_to_corners_format(target_pred_boxes) img_h, img_w = target_sizes.unbind(1) scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) boxes = boxes * scale_fct[:, None, :] iou, _ = box_iou(boxes.squeeze(0), query_box) return iou.argmax() def box_guided_detection( self: OwlViTForObjectDetection, pixel_values, query_box=None, target_sizes=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.return_dict # Compute feature maps for the input and query images # query_feature_map = self.image_embedder(pixel_values=query_pixel_values)[0] feature_map, vision_outputs = self.image_embedder( pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, ) 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)) # batch_size, num_patches, num_patches, hidden_dim = query_feature_map.shape # query_image_feats = torch.reshape(query_feature_map, (batch_size, num_patches * num_patches, hidden_dim)) # # Get top class embedding and best box index for each query image in batch # query_embeds, best_box_indices, query_pred_boxes = self.embed_image_query(query_image_feats, query_feature_map) # Predict object boxes target_pred_boxes = self.box_predictor(image_feats, feature_map) # Get MAX IOU box corresponding embedding query_indice = get_max_iou_indice(target_pred_boxes, query_box, target_sizes) # Predict object classes [batch_size, num_patches, num_queries+1] (pred_logits, class_embeds) = self.class_predictor(image_feats=image_feats, query_indice=query_indice) if not return_dict: output = ( feature_map, # query_feature_map, target_pred_boxes, # query_pred_boxes, pred_logits, class_embeds, vision_outputs.to_tuple(), ) output = tuple(x for x in output if x is not None) return output return OwlViTImageGuidedObjectDetectionOutput( image_embeds=feature_map, # query_image_embeds=query_feature_map, target_pred_boxes=target_pred_boxes, # query_pred_boxes=query_pred_boxes, logits=pred_logits, class_embeds=class_embeds, text_model_output=None, vision_model_output=vision_outputs, ) model.box_guided_detection = partial(box_guided_detection, model) model.class_predictor = partial(class_predictor, model) model.class_head.classpredictionhead_box_forward = partial(classpredictionhead_box_forward, model.class_head) outputs = None def prepare_embedds(xmin, ymin, xmax, ymax, image): box = (int(xmin), int(ymin), int(xmax), int(ymax)) return (image, [(box, "manul")]) def manul_box_change(xmin, ymin, xmax, ymax, image): box = (int(xmin), int(ymin), int(xmax), int(ymax)) return (image["image"], [(box, "manul")]) def threshold_change(xmin, ymin, xmax, ymax, image, threshold, nms): manul_box = (int(xmin), int(ymin), int(xmax), int(ymax)) global outputs target_sizes = torch.Tensor([image["image"].size[::-1]]) results = processor.post_process_image_guided_detection(outputs=outputs, threshold=threshold, nms_threshold=nms, target_sizes=target_sizes) boxes = results[0]['boxes'].type(torch.int64).tolist() scores = results[0]['scores'].tolist() labels = list(zip(boxes, scores)) cnt = len(boxes) return (image["image"], labels), cnt def one_shot_detect(xmin, ymin, xmax, ymax, image, threshold, nms): manul_box = (int(xmin), int(ymin), int(xmax), int(ymax)) global outputs target_sizes = torch.Tensor([image["image"].size[::-1]]) inputs = processor(images=image["image"].convert("RGB"), return_tensors="pt") outputs = model.box_guided_detection(**inputs, query_box=torch.Tensor([manul_box]), target_sizes=target_sizes) results = processor.post_process_image_guided_detection(outputs=outputs, threshold=threshold, nms_threshold=nms, target_sizes=target_sizes) boxes = results[0]['boxes'].type(torch.int64).tolist() scores = results[0]['scores'].tolist() labels = list(zip(boxes, scores)) cnt = len(boxes) return (image["image"], labels), cnt def save_embedding(exam): print(exam) global outputs embedding = outputs["class_embeds"][0, outputs["logits"].argmax()] return embedding.detach().numpy() def sketch2box(sketch_box): mask = sketch_box["mask"].convert("L") mask = np.array(mask) masked_index = np.where(mask == 255) if len(masked_index[0]) == 0: return (sketch_box["image"], []), -1, -1, -1, -1 xmin, ymin, xmax, ymax = masked_index[1].min(), masked_index[0].min(), masked_index[1].max(), masked_index[0].max() box = (xmin, ymin, xmax, ymax) return (sketch_box["image"], [(box, "manual")]), xmin, ymin, xmax, ymax with gr.Blocks() as demo: with gr.Row(): with gr.Column(): sketch_box = gr.Image(type="pil", source="upload", tool="sketch") box_preview = gr.AnnotatedImage(type="pil", interactive=False, height=256) threshold = gr.Number(0.95, label="threshold", step=0.01) nms = gr.Number(0.3, label="nms", step=0.01) cnt = gr.Number(0, label="count", interactive=False) with gr.Column(): annotatedimage = gr.AnnotatedImage() with gr.Row(): xmin = gr.Number(-1, label="xmin") ymin = gr.Number(-1, label="ymin") xmax = gr.Number(-1, label="xmax") ymax = gr.Number(-1, label="ymax") with gr.Row(): run_button = gr.Button(variant="primary") # save_button = gr.Button("Save", variant="secondary") sketch_box.edit(sketch2box, [sketch_box], [box_preview, xmin, ymin, xmax, ymax]) xmin.change(manul_box_change, [xmin, ymin, xmax, ymax, sketch_box], [box_preview]) ymin.change(manul_box_change, [xmin, ymin, xmax, ymax, sketch_box], [box_preview]) xmax.change(manul_box_change, [xmin, ymin, xmax, ymax, sketch_box], [box_preview]) ymax.change(manul_box_change, [xmin, ymin, xmax, ymax, sketch_box], [box_preview]) threshold.change(threshold_change, [xmin, ymin, xmax, ymax, sketch_box, threshold, nms], [annotatedimage, cnt]) nms.change(threshold_change, [xmin, ymin, xmax, ymax, sketch_box, threshold, nms], [annotatedimage, cnt]) run_button.click(one_shot_detect, [xmin, ymin, xmax, ymax, sketch_box, threshold, nms], [annotatedimage, cnt]) demo.launch()