File size: 8,595 Bytes
7a8c937
472aaf0
7a8c937
472aaf0
7a8c937
59ec2c1
7a8c937
 
 
 
472aaf0
7a8c937
59ec2c1
 
0186279
59ec2c1
 
 
472aaf0
7a8c937
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
472aaf0
7a8c937
 
 
 
 
 
 
 
472aaf0
7a8c937
 
 
472aaf0
7a8c937
 
472aaf0
7a8c937
472aaf0
7a8c937
 
 
 
472aaf0
7a8c937
472aaf0
3f492a7
0186279
 
 
 
 
3f492a7
 
0186279
3f492a7
 
 
 
0186279
3f492a7
 
 
 
0186279
3f492a7
0186279
 
3f492a7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7a8c937
 
 
 
472aaf0
 
7a8c937
 
472aaf0
7a8c937
 
472aaf0
 
 
7a8c937
 
 
472aaf0
7a8c937
472aaf0
 
 
7a8c937
472aaf0
7a8c937
 
 
472aaf0
7a8c937
472aaf0
7a8c937
4f3310e
 
 
 
 
7a8c937
472aaf0
 
7a8c937
472aaf0
 
 
 
 
 
 
 
 
1dc1aec
7a8c937
472aaf0
 
 
 
 
7a8c937
472aaf0
 
 
 
7a8c937
472aaf0
 
 
7a8c937
 
 
 
472aaf0
7a8c937
472aaf0
 
 
 
7a8c937
472aaf0
7a8c937
472aaf0
 
4f3310e
472aaf0
 
 
 
4f3310e
 
 
 
 
7a8c937
472aaf0
7a8c937
 
 
472aaf0
 
7a8c937
 
 
472aaf0
 
7a8c937
 
 
 
472aaf0
7a8c937
 
472aaf0
7a8c937
 
472aaf0
7a8c937
 
 
 
 
472aaf0
7a8c937
472aaf0
7a8c937
 
 
 
 
 
 
 
 
 
472aaf0
 
7a8c937
 
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
import gradio as gr
import torch
import numpy as np
from transformers import OwlViTProcessor, OwlViTForObjectDetection
from torchvision import transforms
from PIL import Image, ImageDraw
import cv2
import torch.nn.functional as F
import tempfile
import os
from SuperGluePretrainedNetwork.models.matching import Matching
from SuperGluePretrainedNetwork.models.utils import read_image
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from io import BytesIO

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load models
mixin = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
model = mixin.to(device)

matching = Matching({
    'superpoint': {'nms_radius': 4, 'keypoint_threshold': 0.005, 'max_keypoints': 1024},
    'superglue': {'weights': 'outdoor', 'sinkhorn_iterations': 20, 'match_threshold': 0.2}
}).eval().to(device)

# Utility functions
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return transform(image).unsqueeze(0)

def save_array_to_temp_image(arr):
    rgb_arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(rgb_arr)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
    temp_file_name = temp_file.name
    temp_file.close()
    img.save(temp_file_name)
    return temp_file_name

def stitch_images(images):
    if not images:
        return Image.new('RGB', (100, 100), color='gray')

    max_width = max([img.width for img in images])
    total_height = sum(img.height for img in images)

    composite = Image.new('RGB', (max_width, total_height))

    y_offset = 0
    for img in images:
        composite.paste(img, (0, y_offset))
        y_offset += img.height

    return composite

def unified_matching_plot2(image0, image1, kpts0, kpts1, mkpts0, mkpts1, color, text, path=None, show_keypoints=False, fast_viz=False, opencv_display=False, opencv_title='matches', small_text=[]):
    # Resize images to have the same height
    height = min(image0.shape[0], image1.shape[0])
    image0_resized = cv2.resize(image0, (int(image0.shape[1] * height / image0.shape[0]), height))
    image1_resized = cv2.resize(image1, (int(image1.shape[1] * height / image1.shape[0]), height))

    plt.figure(figsize=(15, 15))
    plt.subplot(1, 2, 1)
    plt.imshow(image0_resized)
    plt.scatter(kpts0[:, 0], kpts0[:, 1], color='r', s=1)
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.imshow(image1_resized)
    plt.scatter(kpts1[:, 0], kpts1[:, 1], color='r', s=1)
    plt.axis('off')

    fig, ax = plt.subplots(figsize=(20, 20))
    plt.plot([mkpts0[:, 0], mkpts1[:, 0] + image0_resized.shape[1]], [mkpts0[:, 1], mkpts1[:, 1]], 'r', lw=0.5)
    plt.scatter(mkpts0[:, 0], mkpts0[:, 1], s=2, marker='o', color='b')
    plt.scatter(mkpts1[:, 0] + image0_resized.shape[1], mkpts1[:, 1], s=2, marker='o', color='g')
    plt.imshow(np.hstack([image0_resized, image1_resized]), aspect='auto')

    plt.suptitle('\n'.join(text), fontsize=20, fontweight='bold')
    plt.tight_layout()
    plt.show()

    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()
    img = cv2.imdecode(img_arr, 1)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.close(fig)

    return img

# Main functions
def detect_and_crop(target_image, query_image, threshold=0.5, nms_threshold=0.3):
    target_sizes = torch.Tensor([target_image.size[::-1]])
    inputs = processor(images=target_image, query_images=query_image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model.image_guided_detection(**inputs)

    img = cv2.cvtColor(np.array(target_image), cv2.COLOR_BGR2RGB)
    outputs.logits = outputs.logits.cpu()
    outputs.target_pred_boxes = outputs.target_pred_boxes.cpu()

    results = processor.post_process_image_guided_detection(outputs=outputs, threshold=threshold, nms_threshold=nms_threshold, target_sizes=target_sizes)
    boxes, scores = results[0]["boxes"], results[0]["scores"]

    if len(boxes) == 0:
        return [], None

    filtered_boxes = []
    for box in boxes:
        x1, y1, x2, y2 = [int(i) for i in box.tolist()]
        cropped_img = img[y1:y2, x1:x2]
        if cropped_img.size != 0:
            filtered_boxes.append(cropped_img)

    draw = ImageDraw.Draw(target_image)
    for box in boxes:
        draw.rectangle(box.tolist(), outline="red", width=3)

    return filtered_boxes, target_image

def image_matching_no_pyramid(query_img, target_img, visualize=True):
    temp_query = save_array_to_temp_image(np.array(query_img))
    temp_target = save_array_to_temp_image(np.array(target_img))
    
    image1, inp1, scales1 = read_image(temp_target, device, [640*2], 0, True)
    image0, inp0, scales0 = read_image(temp_query, device, [640*2], 0, True)

    if image0 is None or image1 is None:
        return None

    pred = matching({'image0': inp0, 'image1': inp1})
    pred = {k: v[0] for k, v in pred.items()}
    kpts0, kpts1 = pred['keypoints0'], pred['keypoints1']
    matches, conf = pred['matches0'], pred['matching_scores0']

    valid = matches > -1
    mkpts0 = kpts0[valid]
    mkpts1 = kpts1[matches[valid]]
    mconf = conf[valid]
    color = cm.jet(mconf.detach().cpu().numpy())[:len(mkpts0)]

    valid_count = np.sum(valid.tolist())

    mkpts0_np = mkpts0.cpu().numpy()
    mkpts1_np = mkpts1.cpu().numpy()

    try:
        H, inliers = cv2.findHomography(mkpts0_np, mkpts1_np, cv2.RANSAC, 5.0)
    except:
        inliers = 0

    num_inliers = np.sum(inliers)

    if visualize:
        visualized_img = unified_matching_plot2(
            image0, image1, kpts0, kpts1, mkpts0, mkpts1, color, ['Matches'], True, False, True, 'Matches', [])
    else:
        visualized_img = None

    return {
        'valid': [valid_count],
        'inliers': [num_inliers],
        'visualized_image': [visualized_img]
    }

def check_object_in_image(query_image, target_image, threshold=50, scale_factor=[0.33, 0.66, 1]):
    images_to_return = []
    cropped_images, bbox_image = detect_and_crop(target_image, query_image)

    temp_files = [save_array_to_temp_image(i) for i in cropped_images]
    crop_results = [image_matching_no_pyramid(query_image, Image.open(i), visualize=True) for i in temp_files]

    cropped_visuals = []
    cropped_inliers = []
    for result in crop_results:
        if result:
            for img in result['visualized_image']:
                cropped_visuals.append(Image.fromarray(img))
            for inliers_ in result['inliers']:
                cropped_inliers.append(inliers_)

    images_to_return.append(stitch_images(cropped_visuals))

    is_present = any(value >= threshold for value in cropped_inliers)

    return {
        'is_present': is_present,
        'images': images_to_return,
        'object detection inliers': [int(i) for i in cropped_inliers],
        'bbox_image': bbox_image,
    }

def interface(poster_source, media_source, threshold, scale_factor):
    result1 = check_object_in_image(poster_source, media_source, threshold, scale_factor)
    if result1['is_present']:
        return result1

    result2 = check_object_in_image(poster_source, media_source, threshold, scale_factor)
    return result2 if result2['is_present'] else result1

iface = gr.Interface(
    fn=interface,
    inputs=[
        gr.Image(type="pil", label="Upload a Query Image (Poster)"),
        gr.Image(type="pil", label="Upload a Target Image (Media)"),
        gr.Slider(minimum=0, maximum=100, step=1, value=50, label="Threshold"),
        gr.CheckboxGroup(choices=[0.33, 0.66, 1.0], value=[0.33, 0.66, 1.0], label="Scale Factors")
    ],
    outputs=[
        gr.JSON(label="Result")
    ],
    title="Object Detection in Image",
    description="""
    **Instructions:**

    1. **Upload a Query Image (Poster)**: Select an image file that contains the object you want to detect.
    2. **Upload a Target Image (Media)**: Select an image file where you want to detect the object.
    3. **Set Threshold**: Adjust the slider to set the threshold for object detection.
    4. **Set Scale Factors**: Select the scale factors for image pyramid.
    5. **View Results**: The result will show whether the object is present in the image along with additional details.
    """
)

if __name__ == "__main__":
    iface.launch()