|
|
from typing import List, Union |
|
|
from PIL import Image, ImageDraw |
|
|
import torch |
|
|
|
|
|
from diffusers.modular_pipelines import ( |
|
|
PipelineState, |
|
|
ModularPipelineBlocks, |
|
|
InputParam, |
|
|
ComponentSpec, |
|
|
OutputParam, |
|
|
) |
|
|
from transformers import AutoProcessor, AutoModelForCausalLM |
|
|
|
|
|
|
|
|
class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): |
|
|
@property |
|
|
def expected_components(self): |
|
|
return [ |
|
|
ComponentSpec( |
|
|
name="image_annotator", |
|
|
type_hint=AutoModelForCausalLM, |
|
|
repo="microsoft/Florence-2-large", |
|
|
), |
|
|
ComponentSpec( |
|
|
name="image_annotator_processor", |
|
|
type_hint=AutoProcessor, |
|
|
repo="microsoft/Florence-2-large", |
|
|
), |
|
|
] |
|
|
|
|
|
@property |
|
|
def inputs(self) -> List[InputParam]: |
|
|
return [ |
|
|
InputParam( |
|
|
"image", |
|
|
Image, |
|
|
required=True, |
|
|
description="Image(s) to annotate", |
|
|
), |
|
|
InputParam( |
|
|
"annotation_task_prompt", |
|
|
Union[str, List[str]], |
|
|
required=True, |
|
|
description="""Annotation Task to perform on the image. |
|
|
""", |
|
|
), |
|
|
] |
|
|
|
|
|
@property |
|
|
def intermediates_outputs(self) -> List[OutputParam]: |
|
|
return [ |
|
|
OutputParam( |
|
|
"mask", |
|
|
type_hint=torch.Tensor, |
|
|
description="Depth Map(s) of input Image(s)", |
|
|
), |
|
|
] |
|
|
|
|
|
def annotate_image(self, image, prompt): |
|
|
inputs = self.image_annotator_processor( |
|
|
text=prompt, images=image, return_tensors="pt" |
|
|
) |
|
|
generated_ids = self.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 = self.image_annotator_processor.batch_decode( |
|
|
generated_ids, skip_special_tokens=False |
|
|
)[0] |
|
|
annotations = self.image_annotator_processor.post_process_generation( |
|
|
annotations, task=prompt, image_size=(image.height, image.width) |
|
|
) |
|
|
|
|
|
return annotations |
|
|
|
|
|
def prepare_mask(self, images, annotations): |
|
|
masks = [] |
|
|
for image, annotation in zip(images, annotations): |
|
|
mask_image = Image.new("L", image.size, 0) |
|
|
draw = ImageDraw.Draw(mask_image) |
|
|
draw.polygon(annotation["polygon"], fill="white") |
|
|
masks.append(mask_image) |
|
|
|
|
|
return masks |
|
|
|
|
|
@torch.no_grad() |
|
|
def __call__(self, pipeline, state: PipelineState) -> PipelineState: |
|
|
block_state = self.get_block_state(state) |
|
|
|
|
|
images = block_state.image |
|
|
annotation_task_prompt = block_state.annotation_task_prompt |
|
|
|
|
|
if not isinstance(annotation_task_prompt, list): |
|
|
annotation_task_prompt = [annotation_task_prompt] |
|
|
|
|
|
if len(images) != len(annotation_task_prompt): |
|
|
raise ValueError("Number of images and annotation prompts must match") |
|
|
|
|
|
annotations = self.annotate_image(images, annotation_task_prompt) |
|
|
block_state.mask = self.prepare_mask(images, annotations) |
|
|
|
|
|
self.set_block_state(block_state) |
|
|
|
