File size: 10,503 Bytes
0041362
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
import ast
import os
import json
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
import matplotlib.pyplot as plt
import numpy as np
import cv2
import inflect

p = inflect.engine()

img_dir = "imgs"
bg_prompt_text = "Background prompt: "
# h, w
box_scale = (512, 512)
size = box_scale
size_h, size_w = size
print(f"Using box scale: {box_scale}")

def parse_input(text=None, no_input=False):
    if not text:
        if no_input:
            return
        
        text = input("Enter the response: ")
    if "Objects: " in text:
        text = text.split("Objects: ")[1]
        
    text_split = text.split(bg_prompt_text)
    if len(text_split) == 2:
        gen_boxes, bg_prompt = text_split
    elif len(text_split) == 1:
        if no_input:
            return
        gen_boxes = text
        bg_prompt = ""
        while not bg_prompt:
            # Ignore the empty lines in the response
            bg_prompt = input("Enter the background prompt: ").strip()
        if bg_prompt_text in bg_prompt:
            bg_prompt = bg_prompt.split(bg_prompt_text)[1]
    else:
        raise ValueError(f"text: {text}")
    try:
        gen_boxes = ast.literal_eval(gen_boxes)    
    except SyntaxError as e:
        # Sometimes the response is in plain text
        if "No objects" in gen_boxes:
            gen_boxes = []
        else:
            raise e
    bg_prompt = bg_prompt.strip()
    
    return gen_boxes, bg_prompt

def filter_boxes(gen_boxes, scale_boxes=True, ignore_background=True, max_scale=3):
    if len(gen_boxes) == 0:
        return []
    
    box_dict_format = False
    gen_boxes_new = []
    for gen_box in gen_boxes:
        if isinstance(gen_box, dict):
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box['name'], gen_box['bounding_box']
            box_dict_format = True
        else:
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box
        if bbox_w <= 0 or bbox_h <= 0:
            # Empty boxes
            continue
        if ignore_background:
            if (bbox_w >= size[1] and bbox_h >= size[0]) or bbox_x > size[1] or bbox_y > size[0]:
                # Ignore the background boxes
                continue
        gen_boxes_new.append(gen_box)
    
    gen_boxes = gen_boxes_new
    
    if len(gen_boxes) == 0:
        return []
    
    filtered_gen_boxes = []
    if box_dict_format:
        # For compatibility
        bbox_left_x_min = min([gen_box['bounding_box'][0] for gen_box in gen_boxes])
        bbox_right_x_max = max([gen_box['bounding_box'][0] + gen_box['bounding_box'][2] for gen_box in gen_boxes])
        bbox_top_y_min = min([gen_box['bounding_box'][1] for gen_box in gen_boxes])
        bbox_bottom_y_max = max([gen_box['bounding_box'][1] + gen_box['bounding_box'][3] for gen_box in gen_boxes])
    else:
        bbox_left_x_min = min([gen_box[1][0] for gen_box in gen_boxes])
        bbox_right_x_max = max([gen_box[1][0] + gen_box[1][2] for gen_box in gen_boxes])
        bbox_top_y_min = min([gen_box[1][1] for gen_box in gen_boxes])
        bbox_bottom_y_max = max([gen_box[1][1] + gen_box[1][3] for gen_box in gen_boxes])
    
    # All boxes are empty
    if (bbox_right_x_max - bbox_left_x_min) == 0:
        return []
    
    # Used if scale_boxes is True
    shift = -bbox_left_x_min
    scale = size_w / (bbox_right_x_max - bbox_left_x_min)
    
    scale = min(scale, max_scale)
    
    for gen_box in gen_boxes:
        if box_dict_format:
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box['name'], gen_box['bounding_box']
        else:
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box
            
        if scale_boxes:
            # Vertical: move the boxes if out of bound
            # Horizontal: move and scale the boxes so it spans the horizontal line
            
            bbox_x = (bbox_x + shift) * scale
            bbox_y = bbox_y * scale
            bbox_w, bbox_h = bbox_w * scale, bbox_h * scale
            # TODO: verify this makes the y center not moving
            bbox_y_offset = 0
            if bbox_top_y_min * scale + bbox_y_offset < 0:
                bbox_y_offset -= bbox_top_y_min * scale
            if bbox_bottom_y_max * scale + bbox_y_offset >= size_h:
                bbox_y_offset -= bbox_bottom_y_max * scale - size_h
            bbox_y += bbox_y_offset
            
            if bbox_y < 0:
                bbox_y, bbox_h = 0, bbox_h - bbox_y
                
        name = name.rstrip(".")
        bounding_box = (int(np.round(bbox_x)), int(np.round(bbox_y)), int(np.round(bbox_w)), int(np.round(bbox_h)))
        if box_dict_format:
            gen_box = {
                'name': name,
                'bounding_box': bounding_box
            }
        else:
            gen_box = (name, bounding_box)
        
        filtered_gen_boxes.append(gen_box)
        
    return filtered_gen_boxes

def draw_boxes(anns):
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in anns:
        c = (np.random.random((1, 3))*0.6+0.4)
        [bbox_x, bbox_y, bbox_w, bbox_h] = ann['bbox']
        poly = [[bbox_x, bbox_y], [bbox_x, bbox_y+bbox_h],
                [bbox_x+bbox_w, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y]]
        np_poly = np.array(poly).reshape((4, 2))
        polygons.append(Polygon(np_poly))
        color.append(c)

        # print(ann)
        name = ann['name'] if 'name' in ann else str(ann['category_id'])
        ax.text(bbox_x, bbox_y, name, style='italic',
                bbox={'facecolor': 'white', 'alpha': 0.7, 'pad': 5})

    p = PatchCollection(polygons, facecolor='none',
                        edgecolors=color, linewidths=2)
    ax.add_collection(p)


def show_boxes(gen_boxes, bg_prompt=None, ind=None, show=False):
    if len(gen_boxes) == 0:
        return
    
    if isinstance(gen_boxes[0], dict):
        anns = [{'name': gen_box['name'], 'bbox': gen_box['bounding_box']}
                for gen_box in gen_boxes]
    else:
        anns = [{'name': gen_box[0], 'bbox': gen_box[1]} for gen_box in gen_boxes]

    # White background (to allow line to show on the edge)
    I = np.ones((size[0]+4, size[1]+4, 3), dtype=np.uint8) * 255

    plt.imshow(I)
    plt.axis('off')

    if bg_prompt is not None:
        ax = plt.gca()
        ax.text(0, 0, bg_prompt, style='italic',
                bbox={'facecolor': 'white', 'alpha': 0.7, 'pad': 5})

        c = (np.zeros((1, 3)))
        [bbox_x, bbox_y, bbox_w, bbox_h] = (0, 0, size[1], size[0])
        poly = [[bbox_x, bbox_y], [bbox_x, bbox_y+bbox_h],
                [bbox_x+bbox_w, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y]]
        np_poly = np.array(poly).reshape((4, 2))
        polygons = [Polygon(np_poly)]
        color = [c]
        p = PatchCollection(polygons, facecolor='none',
                            edgecolors=color, linewidths=2)
        ax.add_collection(p)

    draw_boxes(anns)
    if show:
        plt.show()
    else:
        print("Saved to", f"{img_dir}/boxes.png", f"ind: {ind}")
        if ind is not None:
            plt.savefig(f"{img_dir}/boxes_{ind}.png")
        plt.savefig(f"{img_dir}/boxes.png")


def show_masks(masks):
    masks_to_show = np.zeros((*size, 3), dtype=np.float32)
    for mask in masks:
        c = (np.random.random((3,))*0.6+0.4)

        masks_to_show += mask[..., None] * c[None, None, :]
    plt.imshow(masks_to_show)
    plt.savefig(f"{img_dir}/masks.png")
    plt.show()
    plt.clf()

def convert_box(box, height, width):
    # box: x, y, w, h (in 512 format) -> x_min, y_min, x_max, y_max
    x_min, y_min = box[0] / width, box[1] / height
    w_box, h_box = box[2] / width, box[3] / height
    
    x_max, y_max = x_min + w_box, y_min + h_box
    
    return x_min, y_min, x_max, y_max

def convert_spec(spec, height, width, include_counts=True, verbose=False):
    # Infer from spec
    prompt, gen_boxes, bg_prompt = spec['prompt'], spec['gen_boxes'], spec['bg_prompt']
    
    # This ensures the same objects appear together because flattened `overall_phrases_bboxes` should EXACTLY correspond to `so_prompt_phrase_box_list`. 
    gen_boxes = sorted(gen_boxes, key=lambda gen_box: gen_box[0])
    
    gen_boxes = [(name, convert_box(box, height=height, width=width)) for name, box in gen_boxes]
    
    # NOTE: so phrase should include all the words associated to the object (otherwise "an orange dog" may be recognized as "an orange" by the model generating the background).
    # so word should have one token that includes the word to transfer cross attention (the object name).
    # Currently using the last word of the object name as word.
    if bg_prompt:
        so_prompt_phrase_word_box_list = [(f"{bg_prompt} with {name}", name, name.split(" ")[-1], box) for name, box in gen_boxes]
    else:
        so_prompt_phrase_word_box_list = [(f"{name}", name, name.split(" ")[-1], box) for name, box in gen_boxes]
    
    objects = [gen_box[0] for gen_box in gen_boxes]
    
    objects_unique, objects_count = np.unique(objects, return_counts=True)

    num_total_matched_boxes = 0
    overall_phrases_words_bboxes = []
    for ind, object_name in enumerate(objects_unique):
        bboxes = [box for name, box in gen_boxes if name == object_name]
        
        if objects_count[ind] > 1:
            phrase = p.plural_noun(object_name.replace("an ", "").replace("a ", ""))
            if include_counts:
                phrase = p.number_to_words(objects_count[ind]) + " " + phrase
        else:
            phrase = object_name
        # Currently using the last word of the phrase as word.
        word = phrase.split(' ')[-1]
        
        num_total_matched_boxes += len(bboxes)
        overall_phrases_words_bboxes.append((phrase, word, bboxes))
        
    assert num_total_matched_boxes == len(gen_boxes), f"{num_total_matched_boxes} != {len(gen_boxes)}"

    objects_str = ", ".join([phrase for phrase, _, _ in overall_phrases_words_bboxes])
    if objects_str:
        if bg_prompt:
            overall_prompt = f"{bg_prompt} with {objects_str}"
        else:
            overall_prompt = objects_str
    else:
        overall_prompt = bg_prompt
        
    if verbose:
        print("so_prompt_phrase_word_box_list:", so_prompt_phrase_word_box_list)
        print("overall_prompt:", overall_prompt)
        print("overall_phrases_words_bboxes:", overall_phrases_words_bboxes)
    
    return so_prompt_phrase_word_box_list, overall_prompt, overall_phrases_words_bboxes