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