Spaces:
Running
Running
| import seaborn as sns | |
| from PIL import Image, ImageDraw, ImageFont | |
| import matplotlib.font_manager | |
| import spacy | |
| import re | |
| nlp = spacy.load("en_core_web_sm-3.6.0") | |
| def draw_boxes(image, boxes, texts, output_fn='output.png'): | |
| box_width = 5 | |
| color_palette = sns.color_palette("husl", len(boxes)) | |
| colors = [(int(r*255), int(g*255), int(b*255)) for r, g, b in color_palette] | |
| width, height = image.size | |
| absolute_boxes = [[(int(box[0] * width), int(box[1] * height), int(box[2] * width), int(box[3] * height)) for box in b] for b in boxes] | |
| overlay = Image.new('RGBA', image.size, (255, 255, 255, 0)) | |
| draw = ImageDraw.Draw(overlay) | |
| font_path = sorted(matplotlib.font_manager.findSystemFonts(fontpaths=None, fontext='ttf'))[0] | |
| font = ImageFont.truetype(font_path, size=26) | |
| for box, text, color in zip(absolute_boxes, texts, colors): | |
| for b in box: | |
| draw.rectangle(b, outline=color, width=box_width) | |
| if not text: | |
| continue | |
| splited_text = text.split('\n') | |
| num_lines = len(splited_text) | |
| text_width, text_height = font.getbbox(splited_text[0])[-2:] | |
| y_start = b[3] - text_height * num_lines - box_width | |
| if b[2] - b[0] < 100 or b[3] - b[1] < 100: | |
| y_start = b[3] | |
| for i, line in enumerate(splited_text): | |
| text_width, text_height = font.getbbox(line)[-2:] | |
| x = b[0] + box_width | |
| y = y_start + text_height * i | |
| draw.rectangle([x, y, x+text_width, y+text_height], fill=(128, 128, 128, 160)) | |
| draw.text((x, y), line, font=font, fill=(255, 255, 255)) | |
| img_with_overlay = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB') | |
| img_with_overlay.save(output_fn) | |
| def boxstr_to_boxes(box_str): | |
| boxes = [[int(y)/1000 for y in x.split(',')] for x in box_str.split(';') if x.replace(',', '').isdigit()] | |
| return boxes | |
| def text_to_dict(text): | |
| doc = nlp(text) | |
| box_matches = list(re.finditer(r'\[\[([^\]]+)\]\]', text)) | |
| box_positions = [match.start() for match in box_matches] | |
| noun_phrases = [] | |
| boxes = [] | |
| for match, box_position in zip(box_matches, box_positions): | |
| nearest_np_start = max([0] + [chunk.start_char for chunk in doc.noun_chunks if chunk.end_char <= box_position]) | |
| noun_phrase = text[nearest_np_start:box_position].strip() | |
| if noun_phrase and noun_phrase[-1] == '?': | |
| noun_phrase = text[:box_position].strip() | |
| box_string = match.group(1) | |
| noun_phrases.append(noun_phrase) | |
| boxes.append(boxstr_to_boxes(box_string)) | |
| pairs = [] | |
| for noun_phrase, box_string in zip(noun_phrases, boxes): | |
| pairs.append((noun_phrase.lower(), box_string)) | |
| return dict(pairs) | |
| def parse_response(img, response, output_fn='output.png'): | |
| img = img.convert('RGB') | |
| width, height = img.size | |
| ratio = min(1920 / width, 1080 / height) | |
| new_width = int(width * ratio) | |
| new_height = int(height * ratio) | |
| new_img = img.resize((new_width, new_height), Image.LANCZOS) | |
| pattern = r"\[\[(.*?)\]\]" | |
| positions = re.findall(pattern, response) | |
| boxes = [[[int(y) for y in x.split(',')] for x in pos.split(';') if x.replace(',', '').isdigit()] for pos in positions] | |
| dic = text_to_dict(response) | |
| if not dic: | |
| texts = [] | |
| boxes = [] | |
| else: | |
| texts, boxes = zip(*dic.items()) | |
| draw_boxes(new_img, boxes, texts, output_fn=output_fn) |