365 lines
19 KiB
Python
365 lines
19 KiB
Python
# Based on code from ComfyUI-Florence2 by Jukka Seppänen (kijai), MIT License
|
|
"""Florence-2 processor: image preprocessing, tokenization, task prompts, and post-processing."""
|
|
|
|
import re
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def preprocess_image(image, size=768, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
|
|
"""[B, C, H, W] or [C, H, W] float [0,1] -> [B, C, size, size] normalized."""
|
|
if image.ndim == 3:
|
|
image = image.unsqueeze(0)
|
|
image = torch.nn.functional.interpolate(image, size=(size, size), mode='bicubic', align_corners=False).clamp(0, 1)
|
|
mean_t = torch.tensor(mean, device=image.device, dtype=image.dtype).view(1, 3, 1, 1)
|
|
std_t = torch.tensor(std, device=image.device, dtype=image.dtype).view(1, 3, 1, 1)
|
|
return (image - mean_t) / std_t
|
|
|
|
|
|
class BoxQuantizer:
|
|
def __init__(self, mode, bins):
|
|
self.mode = mode
|
|
self.bins = bins
|
|
|
|
def dequantize(self, boxes, size):
|
|
bins_w, bins_h = self.bins
|
|
size_w, size_h = size
|
|
xmin, ymin, xmax, ymax = boxes.split(1, dim=-1)
|
|
return torch.cat((
|
|
(xmin + 0.5) * size_w / bins_w, (ymin + 0.5) * size_h / bins_h,
|
|
(xmax + 0.5) * size_w / bins_w, (ymax + 0.5) * size_h / bins_h,
|
|
), dim=-1)
|
|
|
|
|
|
class CoordinatesQuantizer:
|
|
def __init__(self, mode, bins):
|
|
self.mode = mode
|
|
self.bins = bins
|
|
|
|
def dequantize(self, coordinates, size):
|
|
bins_w, bins_h = self.bins
|
|
size_w, size_h = size
|
|
x, y = coordinates.split(1, dim=-1)
|
|
return torch.cat(((x + 0.5) * size_w / bins_w, (y + 0.5) * size_h / bins_h), dim=-1)
|
|
|
|
|
|
class PostProcessor:
|
|
"""Regex-based parsing of Florence-2 text outputs into structured results."""
|
|
|
|
def __init__(self, tokenizer=None):
|
|
config = {
|
|
'NUM_BBOX_HEIGHT_BINS': 1000, 'NUM_BBOX_WIDTH_BINS': 1000, 'BOX_QUANTIZATION_MODE': 'floor',
|
|
'COORDINATES_HEIGHT_BINS': 1000, 'COORDINATES_WIDTH_BINS': 1000, 'COORDINATES_QUANTIZATION_MODE': 'floor',
|
|
'PARSE_TASKS': [
|
|
{'TASK_NAME': 'od', 'PATTERN': r'([a-zA-Z0-9 ]+)<loc_(\\d+)><loc_(\\d+)><loc_(\\d+)><loc_(\\d+)>'},
|
|
{'TASK_NAME': 'ocr', 'PATTERN': r'(.+?)<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>', 'AREA_THRESHOLD': 0.00},
|
|
{'TASK_NAME': 'phrase_grounding', 'FILTER_BY_BLACK_LIST': True},
|
|
{'TASK_NAME': 'pure_text'}, {'TASK_NAME': 'description_with_bboxes'},
|
|
{'TASK_NAME': 'description_with_polygons'}, {'TASK_NAME': 'polygons'},
|
|
{'TASK_NAME': 'bboxes'}, {'TASK_NAME': 'description_with_bboxes_or_polygons'},
|
|
],
|
|
}
|
|
self.config = config
|
|
self.parse_tasks = [t['TASK_NAME'] for t in config['PARSE_TASKS']]
|
|
self.parse_tasks_configs = {t['TASK_NAME']: t for t in config['PARSE_TASKS']}
|
|
|
|
self.tokenizer = tokenizer
|
|
if tokenizer is not None:
|
|
self.all_special_tokens = set(tokenizer.all_special_tokens)
|
|
|
|
self.box_quantizer = BoxQuantizer('floor', (1000, 1000))
|
|
self.coordinates_quantizer = CoordinatesQuantizer('floor', (1000, 1000))
|
|
|
|
self.black_list_of_phrase_grounding = set()
|
|
if 'phrase_grounding' in self.parse_tasks and self.parse_tasks_configs['phrase_grounding'].get('FILTER_BY_BLACK_LIST'):
|
|
self.black_list_of_phrase_grounding = {
|
|
'it', 'I', 'me', 'mine', 'you', 'your', 'yours', 'he', 'him', 'his',
|
|
'she', 'her', 'hers', 'they', 'them', 'their', 'theirs', 'one', 'oneself',
|
|
'we', 'us', 'our', 'ours', 'mine', 'yours', 'his', 'hers', 'its',
|
|
'ours', 'yours', 'theirs', 'myself', 'yourself', 'himself', 'herself',
|
|
'itself', 'ourselves', 'yourselves', 'themselves', 'this', 'that',
|
|
'these', 'those', 'who', 'whom', 'whose', 'which', 'what', 'that',
|
|
'all', 'another', 'any', 'anybody', 'anyone', 'anything',
|
|
'each', 'everybody', 'everyone', 'everything',
|
|
'few', 'many', 'nobody', 'none', 'one', 'several',
|
|
'some', 'somebody', 'someone', 'something',
|
|
'each other', 'one another',
|
|
'the image', 'image', 'images', 'the', 'a', 'an', 'a group',
|
|
'other objects', 'lots', 'a set',
|
|
}
|
|
|
|
def parse_od_from_text_and_spans(self, text, pattern, image_size, phrase_centric=False):
|
|
parsed = list(re.finditer(pattern, text))
|
|
instances = []
|
|
for m in parsed:
|
|
if phrase_centric:
|
|
bbox_bins = [int(m.group(j)) for j in range(2, 6)]
|
|
cat_name = m.group(1).lower().strip()
|
|
else:
|
|
bbox_bins = [int(m.group(j)) for j in range(1, 5)]
|
|
cat_name = m.group(5).lower().strip()
|
|
instances.append({
|
|
'bbox': self.box_quantizer.dequantize(torch.tensor(bbox_bins), image_size).tolist(),
|
|
'cat_name': cat_name,
|
|
})
|
|
return instances
|
|
|
|
def parse_ocr_from_text_and_spans(self, text, pattern, image_size, area_threshold=-1.0):
|
|
text = text.replace('<s>', '')
|
|
parsed = re.findall(pattern, text)
|
|
instances = []
|
|
image_width, image_height = image_size
|
|
for ocr_line in parsed:
|
|
ocr_content = ocr_line[0]
|
|
quad_box = [int(i) for i in ocr_line[1:]]
|
|
quad_box = self.coordinates_quantizer.dequantize(torch.tensor(np.array(quad_box).reshape(-1, 2)), image_size).reshape(-1).tolist()
|
|
if area_threshold > 0:
|
|
x_coords, y_coords = quad_box[0::2], quad_box[1::2]
|
|
area = 0.5 * abs(sum(x_coords[i] * y_coords[i + 1] - x_coords[i + 1] * y_coords[i] for i in range(3)))
|
|
if area < (image_width * image_height) * area_threshold:
|
|
continue
|
|
instances.append({'quad_box': quad_box, 'text': ocr_content})
|
|
return instances
|
|
|
|
def parse_phrase_grounding_from_text_and_spans(self, text, pattern, image_size):
|
|
text = text.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
|
|
pattern = r"([^<]+(?:<loc_\d+>){4,})"
|
|
phrases = re.findall(pattern, text)
|
|
phrase_pattern = r'^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_)'
|
|
box_pattern = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
|
|
|
|
instances = []
|
|
for pharse_text in phrases:
|
|
phrase_text_strip = pharse_text.replace('<ground>', '', 1).replace('<obj>', '', 1)
|
|
if phrase_text_strip == '':
|
|
continue
|
|
phrase = re.search(phrase_pattern, phrase_text_strip)
|
|
if phrase is None:
|
|
continue
|
|
bboxes_parsed = list(re.finditer(box_pattern, pharse_text))
|
|
if not bboxes_parsed:
|
|
continue
|
|
phrase = phrase.group().strip()
|
|
if phrase in self.black_list_of_phrase_grounding:
|
|
continue
|
|
bbox_bins = [[int(b.group(j)) for j in range(1, 5)] for b in bboxes_parsed]
|
|
phrase = phrase.encode('ascii', errors='ignore').decode('ascii')
|
|
instances.append({
|
|
'bbox': self.box_quantizer.dequantize(torch.tensor(bbox_bins), image_size).tolist(),
|
|
'cat_name': phrase,
|
|
})
|
|
return instances
|
|
|
|
def parse_description_with_bboxes_from_text_and_spans(self, text, pattern, image_size, allow_empty_phrase=False):
|
|
text = text.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
|
|
if allow_empty_phrase:
|
|
pattern = r"(?:(?:<loc_\d+>){{4,}})"
|
|
else:
|
|
pattern = r"([^<]+(?:<loc_\d+>){4,})"
|
|
phrases = re.findall(pattern, text)
|
|
phrase_pattern = r'^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_)'
|
|
box_pattern = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
|
|
|
|
instances = []
|
|
for pharse_text in phrases:
|
|
phrase_text_strip = pharse_text.replace('<ground>', '', 1).replace('<obj>', '', 1)
|
|
if phrase_text_strip == '' and not allow_empty_phrase:
|
|
continue
|
|
phrase = re.search(phrase_pattern, phrase_text_strip)
|
|
if phrase is None:
|
|
continue
|
|
phrase = phrase.group().strip()
|
|
bboxes_parsed = list(re.finditer(box_pattern, pharse_text))
|
|
if not bboxes_parsed:
|
|
continue
|
|
bbox_bins = [[int(b.group(j)) for j in range(1, 5)] for b in bboxes_parsed]
|
|
bboxes = self.box_quantizer.dequantize(torch.tensor(bbox_bins), image_size).tolist()
|
|
phrase = phrase.encode('ascii', errors='ignore').decode('ascii')
|
|
for bbox in bboxes:
|
|
instances.append({'bbox': bbox, 'cat_name': phrase})
|
|
return instances
|
|
|
|
def parse_description_with_polygons_from_text_and_spans(self, text, pattern, image_size,
|
|
allow_empty_phrase=False, polygon_sep_token='<sep>',
|
|
polygon_start_token='<poly>', polygon_end_token='</poly>',
|
|
with_box_at_start=False):
|
|
text = text.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
|
|
sep, start, end = re.escape(polygon_sep_token), re.escape(polygon_start_token), re.escape(polygon_end_token)
|
|
if allow_empty_phrase:
|
|
pattern = rf"(?:(?:<loc_\d+>|{sep}|{start}|{end}){{4,}})"
|
|
else:
|
|
pattern = rf"([^<]+(?:<loc_\d+>|{sep}|{start}|{end}){{4,}})"
|
|
phrases = re.findall(pattern, text)
|
|
phrase_string_pattern = r'^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_|<poly>)'
|
|
box_pattern = rf'((?:<loc_\d+>)+)(?:{sep}|$)'
|
|
polygons_instance_pattern = rf'{start}(.*?){end}'
|
|
|
|
instances = []
|
|
for phrase_text in phrases:
|
|
phrase_text_strip = re.sub(r'^loc_\d+>', '', phrase_text, count=1)
|
|
if phrase_text_strip == '' and not allow_empty_phrase:
|
|
continue
|
|
phrase = re.search(phrase_string_pattern, phrase_text_strip)
|
|
if phrase is None:
|
|
continue
|
|
phrase = phrase.group().strip()
|
|
|
|
if polygon_start_token in phrase_text and polygon_end_token in phrase_text:
|
|
poly_instances = list(re.finditer(polygons_instance_pattern, phrase_text))
|
|
else:
|
|
poly_instances = [phrase_text]
|
|
|
|
for pi in poly_instances:
|
|
poly_text = pi.group(1) if not isinstance(pi, str) else pi
|
|
polygons_parsed = list(re.finditer(box_pattern, poly_text))
|
|
if not polygons_parsed:
|
|
continue
|
|
|
|
bbox, polygons = [], []
|
|
for pp in polygons_parsed:
|
|
coords = [int(m.group(1)) for m in re.finditer(r'<loc_(\d+)>', pp.group(1))]
|
|
if with_box_at_start and not bbox:
|
|
if len(coords) > 4:
|
|
bbox = coords[:4]
|
|
coords = coords[4:]
|
|
else:
|
|
bbox = [0, 0, 0, 0]
|
|
if len(coords) % 2 == 1:
|
|
coords = coords[:-1]
|
|
polygons.append(self.coordinates_quantizer.dequantize(
|
|
torch.tensor(np.array(coords).reshape(-1, 2)), image_size,
|
|
).reshape(-1).tolist())
|
|
|
|
instance = {'cat_name': phrase, 'polygons': polygons}
|
|
if bbox:
|
|
instance['bbox'] = self.box_quantizer.dequantize(torch.tensor([bbox]), image_size).tolist()[0]
|
|
instances.append(instance)
|
|
return instances
|
|
|
|
def __call__(self, text=None, image_size=None, parse_tasks=None):
|
|
if parse_tasks is not None:
|
|
if isinstance(parse_tasks, str):
|
|
parse_tasks = [parse_tasks]
|
|
for t in parse_tasks:
|
|
assert t in self.parse_tasks, f'parse task {t} not supported'
|
|
|
|
assert text is not None, 'text should be provided'
|
|
parsed_dict = {'text': text}
|
|
|
|
for task in self.parse_tasks:
|
|
if parse_tasks is not None and task not in parse_tasks:
|
|
continue
|
|
pattern = self.parse_tasks_configs[task].get('PATTERN', None)
|
|
|
|
if task == 'ocr':
|
|
parsed_dict['ocr'] = self.parse_ocr_from_text_and_spans(text, pattern, image_size, self.parse_tasks_configs[task].get('AREA_THRESHOLD', 0.0))
|
|
elif task == 'phrase_grounding':
|
|
parsed_dict['phrase_grounding'] = self.parse_phrase_grounding_from_text_and_spans(text, pattern, image_size)
|
|
elif task == 'pure_text':
|
|
parsed_dict['pure_text'] = text
|
|
elif task == 'description_with_bboxes':
|
|
parsed_dict['description_with_bboxes'] = self.parse_description_with_bboxes_from_text_and_spans(text, pattern, image_size)
|
|
elif task == 'description_with_polygons':
|
|
parsed_dict['description_with_polygons'] = self.parse_description_with_polygons_from_text_and_spans(text, pattern, image_size)
|
|
elif task == 'polygons':
|
|
parsed_dict['polygons'] = self.parse_description_with_polygons_from_text_and_spans(text, pattern, image_size, allow_empty_phrase=True)
|
|
elif task == 'bboxes':
|
|
parsed_dict['bboxes'] = self.parse_description_with_bboxes_from_text_and_spans(text, pattern, image_size, allow_empty_phrase=True)
|
|
elif task == 'description_with_bboxes_or_polygons':
|
|
if '<poly>' in text:
|
|
parsed_dict['description_with_bboxes_or_polygons'] = self.parse_description_with_polygons_from_text_and_spans(text, pattern, image_size)
|
|
else:
|
|
parsed_dict['description_with_bboxes_or_polygons'] = self.parse_description_with_bboxes_from_text_and_spans(text, pattern, image_size)
|
|
else:
|
|
raise ValueError(f"task {task} is not supported")
|
|
return parsed_dict
|
|
|
|
|
|
class Processor:
|
|
def __init__(self, model_path):
|
|
from .tokenizer import Florence2Tokenizer
|
|
self.tokenizer = Florence2Tokenizer(model_path)
|
|
self.image_seq_length = 577
|
|
|
|
self.tasks_answer_post_processing_type = {
|
|
'<OCR>': 'pure_text', '<OCR_WITH_REGION>': 'ocr',
|
|
'<CAPTION>': 'pure_text', '<DETAILED_CAPTION>': 'pure_text', '<MORE_DETAILED_CAPTION>': 'pure_text',
|
|
'<OD>': 'description_with_bboxes', '<DENSE_REGION_CAPTION>': 'description_with_bboxes',
|
|
'<CAPTION_TO_PHRASE_GROUNDING>': 'phrase_grounding',
|
|
'<REFERRING_EXPRESSION_SEGMENTATION>': 'polygons', '<REGION_TO_SEGMENTATION>': 'polygons',
|
|
'<OPEN_VOCABULARY_DETECTION>': 'description_with_bboxes_or_polygons',
|
|
'<REGION_TO_CATEGORY>': 'pure_text', '<REGION_TO_DESCRIPTION>': 'pure_text', '<REGION_TO_OCR>': 'pure_text',
|
|
'<REGION_PROPOSAL>': 'bboxes',
|
|
}
|
|
self.task_prompts_without_inputs = {
|
|
'<OCR>': 'What is the text in the image?',
|
|
'<OCR_WITH_REGION>': 'What is the text in the image, with regions?',
|
|
'<CAPTION>': 'What does the image describe?',
|
|
'<DETAILED_CAPTION>': 'Describe in detail what is shown in the image.',
|
|
'<MORE_DETAILED_CAPTION>': 'Describe with a paragraph what is shown in the image.',
|
|
'<OD>': 'Locate the objects with category name in the image.',
|
|
'<DENSE_REGION_CAPTION>': 'Locate the objects in the image, with their descriptions.',
|
|
'<REGION_PROPOSAL>': 'Locate the region proposals in the image.',
|
|
}
|
|
self.task_prompts_with_input = {
|
|
'<CAPTION_TO_PHRASE_GROUNDING>': 'Locate the phrases in the caption: {input}',
|
|
'<REFERRING_EXPRESSION_SEGMENTATION>': 'Locate {input} in the image with mask',
|
|
'<REGION_TO_SEGMENTATION>': 'What is the polygon mask of region {input}',
|
|
'<OPEN_VOCABULARY_DETECTION>': 'Locate {input} in the image.',
|
|
'<REGION_TO_CATEGORY>': 'What is the region {input}?',
|
|
'<REGION_TO_DESCRIPTION>': 'What does the region {input} describe?',
|
|
'<REGION_TO_OCR>': 'What text is in the region {input}?',
|
|
}
|
|
self.post_processor = PostProcessor(tokenizer=self.tokenizer)
|
|
|
|
def _construct_prompts(self, text):
|
|
for task_token, task_prompt in self.task_prompts_without_inputs.items():
|
|
if task_token in text:
|
|
return task_prompt
|
|
for task_token, task_prompt in self.task_prompts_with_input.items():
|
|
if task_token in text:
|
|
return task_prompt.format(input=text.replace(task_token, ''))
|
|
return text
|
|
|
|
def __call__(self, text, images):
|
|
prompt = self._construct_prompts(text)
|
|
encoded = self.tokenizer.encode(prompt)
|
|
pixel_values = preprocess_image(images)
|
|
return {'input_ids': encoded['input_ids'], 'pixel_values': pixel_values}
|
|
|
|
def batch_decode(self, token_ids, skip_special_tokens=False):
|
|
return self.tokenizer.batch_decode(token_ids, skip_special_tokens=skip_special_tokens)
|
|
|
|
def post_process_generation(self, text, task, image_size):
|
|
pp_type = self.tasks_answer_post_processing_type.get(task, 'pure_text')
|
|
result = self.post_processor(text=text, image_size=image_size, parse_tasks=pp_type)[pp_type]
|
|
|
|
if pp_type == 'pure_text':
|
|
final = result.replace('<s>', '').replace('</s>', '')
|
|
elif pp_type in ('od', 'description_with_bboxes', 'bboxes'):
|
|
final = {'bboxes': [i['bbox'] for i in result], 'labels': [str(i['cat_name']) for i in result]}
|
|
elif pp_type == 'ocr':
|
|
final = {'quad_boxes': [i['quad_box'] for i in result], 'labels': [str(i['text']) for i in result]}
|
|
elif pp_type == 'phrase_grounding':
|
|
bboxes, labels = [], []
|
|
for phrase in result:
|
|
for bbox in phrase['bbox']:
|
|
bboxes.append(bbox)
|
|
labels.append(phrase['cat_name'])
|
|
final = {'bboxes': bboxes, 'labels': labels}
|
|
elif pp_type in ('description_with_polygons', 'polygons'):
|
|
final = {'polygons': [r['polygons'] for r in result], 'labels': [r['cat_name'] for r in result]}
|
|
elif pp_type == 'description_with_bboxes_or_polygons':
|
|
bboxes, bl, polygons, pl = [], [], [], []
|
|
for r in result:
|
|
if 'polygons' in r:
|
|
polygons.append(r['polygons'])
|
|
pl.append(r['cat_name'])
|
|
else:
|
|
bboxes.append(r['bbox'])
|
|
bl.append(r['cat_name'])
|
|
final = {'bboxes': bboxes, 'bboxes_labels': bl, 'polygons': polygons, 'polygons_labels': pl}
|
|
else:
|
|
raise ValueError(f'Unknown post processing type: {pp_type}')
|
|
return {task: final}
|