File size: 4,119 Bytes
a64bccf
 
 
 
e4dee6a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a64bccf
 
 
 
 
 
 
 
 
e4dee6a
 
 
a64bccf
 
e4dee6a
a64bccf
 
 
 
 
 
e4dee6a
 
 
 
 
a64bccf
 
e4dee6a
 
 
a64bccf
e4dee6a
a64bccf
e4dee6a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a64bccf
 
 
e4dee6a
a64bccf
 
e4dee6a
 
 
 
 
a64bccf
e4dee6a
 
 
 
a64bccf
 
e4dee6a
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import torch
import gradio as gr
from transformers import Owlv2Processor, Owlv2ForObjectDetection
import spaces
import numpy as np
from PIL import Image
import io
import random
from transformers import SamModel, SamProcessor


def apply_colored_masks_on_image(image, masks):
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image.astype('uint8'), 'RGB')

    image_rgba = image.convert("RGBA")

    for i in range(masks.shape[0]):
        mask = masks[i].squeeze().cpu().numpy()
        mask_image = Image.fromarray((mask * 255).astype(np.uint8), 'L')
        color = tuple([random.randint(0, 255) for _ in range(3)] + [128]) 
        colored_mask = Image.new("RGBA", image.size, color)
        colored_mask.putalpha(mask_image)
        image_rgba = Image.alpha_composite(image_rgba, colored_mask)

    return image_rgba


# Use GPU if available
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
model_sam = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
processor_sam = SamProcessor.from_pretrained("facebook/sam-vit-huge")


@spaces.GPU
def query_image(img, text_queries, score_threshold=0.5):
    text_queries = text_queries.split(",")
    size = max(img.shape[:2])
    target_sizes = torch.Tensor([[size, size]])
    inputs = processor(text=text_queries, images=img, return_tensors="pt").to(device)

    with torch.no_grad():
        model_outputs = model(**inputs)
        model_outputs.logits = model_outputs.logits.cpu()
        model_outputs.pred_boxes = model_outputs.pred_boxes.cpu()
        results = processor.post_process_object_detection(outputs=model_outputs, target_sizes=target_sizes)
    
    boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]

  
    img_pil = Image.fromarray(img.astype('uint8'), 'RGB')

    result_labels = []
    result_boxes = []
    for box, score, label in zip(boxes, scores, labels):
        if score >= score_threshold:
            box = [int(i) for i in box.tolist()]
            label_text = text_queries[label.item()]
            result_labels.append((box, label_text))
            result_boxes.append(box)

    input_boxes_for_sam = [result_boxes]
    sam_image = generate_image_with_sam(np.array(img_pil), input_boxes_for_sam)

    return sam_image,result_labels

   
def generate_image_with_sam(img, boxes):
    img_pil = Image.fromarray(img.astype('uint8'), 'RGB')
    inputs = processor_sam(img_pil, return_tensors="pt").to(device)
    
    image_embeddings = model_sam.get_image_embeddings(inputs["pixel_values"])

    inputs = processor_sam(img_pil, input_boxes=boxes, return_tensors="pt").to(device)
    inputs["input_boxes"].shape
    inputs.pop("pixel_values", None)
    inputs.update({"image_embeddings": image_embeddings})

    with torch.no_grad():
        outputs = model_sam(**inputs, multimask_output=False)
    
    masks = processor_sam.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
    scores = outputs.iou_scores
    print(type(scores))
    print(scores.shape if hasattr(scores, 'shape') else scores)


    SAM_image = apply_colored_masks_on_image(img_pil, masks[0])
    return SAM_image


description = """
Split anythings
"""
demo = gr.Interface(
    fn=query_image,
    inputs=[gr.Image(), gr.Textbox(label="Query Text"), gr.Slider(0, 1, value=0.5, label="Score Threshold")],
    outputs=gr.AnnotatedImage(),
    title="Zero-Shot Object Detection SV3",
    description="This interface demonstrates object detection using  zero-shot object detection and SAM  for image segmentation.",
    examples=[
        ["images/purple cell.png", "purple cells", 0.11],
        ["images/dark_cell.png", "gray cells", 0.1],
        ["images/animals.png", "Rabbit,Squirrel,Parrot,Hedgehog,Turtle,Ladybug,Chick,Frog,Butterfly,Snail,Mouse", 0.1],

    ],
)

demo.launch()