Update app.py

app.py

@@ -1,16 +1,20 @@
 import gradio as gr
 from PIL import Image, ImageDraw
 import torch
-from transformers import OwlViTProcessor, OwlViTForObjectDetection
+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
+from transformers.models.owlvit.modeling_owlvit import OwlViTImageGuidedObjectDetectionOutput, OwlViTClassPredictionHead
+
+
 
 
 
@@ -24,9 +28,7 @@ def classpredictionhead_box_forward(
 
     # Normalize image and text features
     image_class_embeds = image_class_embeds / (torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6)
-    print(image_class_embeds.shape)
     query_embeds = image_class_embeds[0, query_indice].unsqueeze(0).unsqueeze(0)
-    print(query_embeds.shape)
     # query_embeds = query_embeds / (torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6)
 
     # Get class predictions
@@ -66,6 +68,8 @@ def class_predictor(
 
 
 
+
+
 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)
@@ -104,6 +108,12 @@ def box_guided_detection(
     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
@@ -113,6 +123,9 @@ def box_guided_detection(
     (pred_logits, class_embeds) = self.class_predictor(image_feats=image_feats, query_indice=query_indice)
 
 
+
+
+
     if not return_dict:
         output = (
             feature_map,
@@ -150,31 +163,30 @@ def prepare_embedds(xmin, ymin, xmax, ymax, image):
 
 def manul_box_change(xmin, ymin, xmax, ymax, image):
     box = (int(xmin), int(ymin), int(xmax), int(ymax))
-    return (image, [(box, "manul")])
+    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.size[::-1]])
+    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))
-    labels.append((manul_box, "manual"))
 
     cnt = len(boxes)
 
-    return (image, labels), cnt
+    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.size[::-1]])
-    inputs = processor(images=image.convert("RGB"), return_tensors="pt")
+    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)
@@ -182,37 +194,60 @@ def one_shot_detect(xmin, ymin, xmax, ymax, image, threshold, nms):
     boxes = results[0]['boxes'].type(torch.int64).tolist()
     scores = results[0]['scores'].tolist()
     labels = list(zip(boxes, scores))
-    labels.append((manul_box, "manual"))
 
     cnt = len(boxes)
 
-    return (image, labels), cnt
+    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():
-            image = gr.Image(type="pil")
+            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(8, label="xmin")
-        ymin = gr.Number(198, label="ymin")
-        xmax = gr.Number(100, label="xmax")
-        ymax = gr.Number(428, label="ymax")
-    button = gr.Button(variant="primary")
-
-    xmin.change(manul_box_change, [xmin, ymin, xmax, ymax, image], [annotatedimage])
-    ymin.change(manul_box_change, [xmin, ymin, xmax, ymax, image], [annotatedimage])
-    xmax.change(manul_box_change, [xmin, ymin, xmax, ymax, image], [annotatedimage])
-    ymax.change(manul_box_change, [xmin, ymin, xmax, ymax, image], [annotatedimage])
-    threshold.change(threshold_change, [xmin, ymin, xmax, ymax, image, threshold, nms], [annotatedimage, cnt])
-    nms.change(threshold_change, [xmin, ymin, xmax, ymax, image, threshold, nms], [annotatedimage, cnt])
-    image.upload(prepare_embedds, [xmin, ymin, xmax, ymax, image], [annotatedimage])
-    button.click(one_shot_detect, [xmin, ymin, xmax, ymax, image, threshold, nms], [annotatedimage, cnt])
+        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])
+