| | from typing import List, Union |
| |
|
| | import numpy as np |
| | import torch |
| | from diffusers.modular_pipelines import ( |
| | ComponentSpec, |
| | InputParam, |
| | ModularPipelineBlocks, |
| | OutputParam, |
| | PipelineState, |
| | ) |
| | from PIL import Image, ImageDraw |
| | from transformers import Florence2ForConditionalGeneration, AutoProcessor |
| |
|
| |
|
| | class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): |
| | @property |
| | def expected_components(self): |
| | return [ |
| | ComponentSpec( |
| | name="image_annotator", |
| | type_hint=Florence2ForConditionalGeneration, |
| | repo="florence-community/Florence-2-base-ft", |
| | ), |
| | ComponentSpec( |
| | name="image_annotator_processor", |
| | type_hint=AutoProcessor, |
| | repo="florence-community/Florence-2-base-ft", |
| | ), |
| | ] |
| |
|
| | @property |
| | def inputs(self) -> List[InputParam]: |
| | return [ |
| | InputParam( |
| | "image", |
| | type_hint=Union[Image.Image, List[Image.Image]], |
| | required=True, |
| | description="Image(s) to annotate", |
| | ), |
| | InputParam( |
| | "annotation_task", |
| | type_hint=Union[str, List[str]], |
| | required=True, |
| | default="<REFERRING_EXPRESSION_SEGMENTATION>", |
| | description="""Annotation Task to perform on the image. |
| | Supported Tasks: |
| | |
| | <OD> |
| | <REFERRING_EXPRESSION_SEGMENTATION> |
| | <CAPTION> |
| | <DETAILED_CAPTION> |
| | <MORE_DETAILED_CAPTION> |
| | <DENSE_REGION_CAPTION> |
| | <CAPTION_TO_PHRASE_GROUNDING> |
| | <OPEN_VOCABULARY_DETECTION> |
| | |
| | """, |
| | ), |
| | InputParam( |
| | "annotation_prompt", |
| | type_hint=Union[str, List[str]], |
| | required=True, |
| | description="""Annotation Prompt to provide more context to the task. |
| | Can be used to detect or segment out specific elements in the image |
| | """, |
| | ), |
| | InputParam( |
| | "annotation_output_type", |
| | type_hint=str, |
| | required=True, |
| | default="mask_image", |
| | description="""Output type from annotation predictions. Availabe options are |
| | annotation: |
| | - raw annotation predictions from the model based on task type. |
| | mask_image: |
| | -black and white mask image for the given image based on the task type |
| | mask_overlay: |
| | - white mask overlayed on the original image |
| | bounding_box: |
| | - bounding boxes drawn on the original image |
| | """, |
| | ), |
| | InputParam( |
| | "annotation_overlay", |
| | type_hint=bool, |
| | required=True, |
| | default=False, |
| | description="", |
| | ), |
| | InputParam( |
| | "fill", |
| | type_hint=str, |
| | default="white", |
| | description="", |
| | ), |
| | ] |
| |
|
| | @property |
| | def intermediate_outputs(self) -> List[OutputParam]: |
| | return [ |
| | OutputParam( |
| | "mask_image", |
| | type_hint=Image, |
| | description="Inpainting Mask for input Image(s)", |
| | ), |
| | OutputParam( |
| | "annotations", |
| | type_hint=dict, |
| | description="Annotations Predictions for input Image(s)", |
| | ), |
| | OutputParam( |
| | "image", |
| | type_hint=Image, |
| | description="Annotated input Image(s)", |
| | ), |
| | ] |
| |
|
| | def get_annotations(self, components, images, prompts, task): |
| | task_prompts = [task + prompt for prompt in prompts] |
| |
|
| | inputs = components.image_annotator_processor( |
| | text=task_prompts, images=images, return_tensors="pt" |
| | ).to(components.image_annotator.device, components.image_annotator.dtype) |
| |
|
| | generated_ids = components.image_annotator.generate( |
| | input_ids=inputs["input_ids"], |
| | pixel_values=inputs["pixel_values"], |
| | max_new_tokens=1024, |
| | early_stopping=False, |
| | do_sample=False, |
| | num_beams=3, |
| | ) |
| | annotations = components.image_annotator_processor.batch_decode( |
| | generated_ids, skip_special_tokens=False |
| | ) |
| | outputs = [] |
| | for image, annotation in zip(images, annotations): |
| | outputs.append( |
| | components.image_annotator_processor.post_process_generation( |
| | annotation, task=task, image_size=(image.width, image.height) |
| | ) |
| | ) |
| | return outputs |
| |
|
| | def prepare_mask(self, images, annotations, overlay=False, fill="white"): |
| | masks = [] |
| | for image, annotation in zip(images, annotations): |
| | mask_image = image.copy() if overlay else Image.new("L", image.size, 0) |
| | draw = ImageDraw.Draw(mask_image) |
| |
|
| | for _, _annotation in annotation.items(): |
| | if "polygons" in _annotation: |
| | for polygon in _annotation["polygons"]: |
| | polygon = np.array(polygon).reshape(-1, 2) |
| | if len(polygon) < 3: |
| | continue |
| | polygon = polygon.reshape(-1).tolist() |
| | draw.polygon(polygon, fill=fill) |
| |
|
| | elif "bbox" in _annotation: |
| | bbox = _annotation["bbox"] |
| | draw.rectangle(bbox, fill="white") |
| |
|
| | masks.append(mask_image) |
| |
|
| | return masks |
| |
|
| | def prepare_bounding_boxes(self, images, annotations): |
| | outputs = [] |
| | for image, annotation in zip(images, annotations): |
| | image_copy = image.copy() |
| | draw = ImageDraw.Draw(image_copy) |
| | for _, _annotation in annotation.items(): |
| | bbox = _annotation["bbox"] |
| | label = _annotation["label"] |
| |
|
| | draw.rectangle(bbox, outline="red", width=3) |
| | draw.text((bbox[0], bbox[1] - 20), label, fill="red") |
| |
|
| | outputs.append(image_copy) |
| |
|
| | return outputs |
| |
|
| | def prepare_inputs(self, images, prompts): |
| | prompts = prompts or "" |
| |
|
| | if isinstance(images, Image.Image): |
| | images = [images] |
| | if isinstance(prompts, str): |
| | prompts = [prompts] |
| |
|
| | if len(images) != len(prompts): |
| | raise ValueError("Number of images and annotation prompts must match.") |
| |
|
| | return images, prompts |
| |
|
| | @torch.no_grad() |
| | def __call__(self, components, state: PipelineState) -> PipelineState: |
| | block_state = self.get_block_state(state) |
| | images, annotation_task_prompt = self.prepare_inputs( |
| | block_state.image, block_state.annotation_prompt |
| | ) |
| | task = block_state.annotation_task |
| | fill = block_state.fill |
| |
|
| | annotations = self.get_annotations( |
| | components, images, annotation_task_prompt, task |
| | ) |
| | block_state.annotations = annotations |
| | if block_state.annotation_output_type == "mask_image": |
| | block_state.mask_image = self.prepare_mask(images, annotations) |
| | else: |
| | block_state.mask_image = None |
| |
|
| | if block_state.annotation_output_type == "mask_overlay": |
| | block_state.image = self.prepare_mask( |
| | images, annotations, overlay=True, fill=fill |
| | ) |
| |
|
| | elif block_state.annotation_output_type == "bounding_box": |
| | block_state.image = self.prepare_bounding_boxes(images, annotations) |
| |
|
| | self.set_block_state(state, block_state) |
| |
|
| | return components, state |
| |
|
