Spaces:
Sleeping
Sleeping
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 | |