| """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} |
|
|
