| from typing import Optional, Dict, Union, Tuple, List |
| from PIL import Image |
| import mmengine.fileio as fileio |
| from mmengine.logging import print_log |
| import io |
| from mmcv.transforms import LoadImageFromFile, BaseTransform |
| from xtuner.registry import BUILDER |
| from xtuner.utils.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN |
| import torch |
| import torch.nn.functional as F |
| import copy |
|
|
|
|
| class PILLoadImageFromFile(LoadImageFromFile): |
| def __init__(self, **kwargs): |
| backend_args = kwargs.pop('backend_args', None) |
| super().__init__(backend_args=backend_args, **kwargs) |
|
|
| def transform(self, results: dict) -> Optional[dict]: |
| """Functions to load image. |
| |
| Args: |
| results (dict): Result dict from |
| :class:`mmengine.dataset.BaseDataset`. |
| |
| Returns: |
| dict: The dict contains loaded image and meta information. |
| """ |
|
|
| filename = results['img_path'] |
| try: |
| if self.file_client_args is not None: |
| file_client = fileio.FileClient.infer_client( |
| self.file_client_args, filename) |
| img_bytes = file_client.get(filename) |
| else: |
| img_bytes = fileio.get( |
| filename, backend_args=self.backend_args) |
| img = Image.open(io.BytesIO(img_bytes)) |
| except Exception as e: |
| if self.ignore_empty: |
| return None |
| else: |
| raise e |
| |
| |
| assert img is not None, f'failed to load image: {filename}' |
| results['img'] = img |
| results['img_shape'] = (img.height, img.width) |
| results['ori_shape'] = (img.height, img.width) |
| return results |
|
|
|
|
| class RefCOCO2PNG(BaseTransform): |
| def __init__(self, |
| image_processor=None, |
| tokenizer=None, |
| prompt_template=None, |
| prompt='<image>\nWhat is shown in this image?', |
| concat=True, |
| image2tensor=True, |
| add_image_token=False, |
| image_token=DEFAULT_IMAGE_TOKEN): |
| self.tokenizer = BUILDER.build(tokenizer) |
| self.image_processor = BUILDER.build(image_processor) |
| self.concat = concat |
| self.image2tensor = image2tensor |
| self.image_token = image_token |
|
|
| self.add_image_token = add_image_token |
| if add_image_token: |
| print_log(f"Manually add image token: {self.image_token}") |
| special_tokens_dict = {'additional_special_tokens': [self.image_token, ]} |
| num_added_toks = self.tokenizer.add_special_tokens(special_tokens_dict) |
| assert num_added_toks == 1 |
|
|
| self.image_token_idx = self.tokenizer.encode(self.image_token, add_special_tokens=False)[-1] |
| print_log(f"Image token: {self.tokenizer.decode(self.image_token_idx)}") |
|
|
| self.prompt = self.tokenizer.encode( |
| prompt_template['INSTRUCTION'].format(input=prompt), |
| add_special_tokens=True) |
| self.prompt_template = prompt_template |
|
|
| def transform(self, results): |
| if self.concat: |
| return self.transform_concat(results) |
| else: |
| return self.transform_split(results) |
|
|
| def transform_split(self, results): |
| all_results = [] |
| for inst_id, instant_text in enumerate(results['text']): |
| new_results = copy.deepcopy(results) |
| new_results['text'] = [instant_text] |
| new_results['gt_masks'] = results['gt_masks'][inst_id:inst_id+1] |
| all_results.append(self.transform_concat(new_results)) |
|
|
| return all_results |
|
|
| def transform_concat(self, results: dict): |
|
|
| caption_input_ids = [] |
| mask_ids = [-1] * len(self.prompt) |
| split_token_id = self.tokenizer.encode('.', add_special_tokens=False)[-1] |
|
|
| for inst_id, instant_text in enumerate(results['text']): |
| segment_input_ids = self.tokenizer.encode(instant_text, add_special_tokens=False) |
| caption_input_ids += segment_input_ids |
| mask_ids += [inst_id] * len(segment_input_ids) |
|
|
| caption_input_ids.append(split_token_id) |
| mask_ids.append(-1) |
|
|
| input_ids = self.prompt + caption_input_ids |
| input_ids = torch.tensor(input_ids, dtype=torch.long) |
| mask_ids = torch.tensor(mask_ids) |
|
|
| image = results['img'] |
| image_data = self.image_processor.preprocess(image) |
|
|
| pixel_values = image_data['pixel_values'][0] |
| if self.image2tensor: |
| pixel_values = torch.from_numpy(pixel_values) |
| meta_data = image_data['meta_datas'][0] |
|
|
| assert len(results['gt_masks'].masks) == len(results['text']) |
| mask_cnt = len(results['text']) |
|
|
| masks = torch.from_numpy(results['gt_masks'].masks).float() |
|
|
| h, w = meta_data['image_shape']['height'], meta_data['image_shape']['width'] |
| gt_masks = masks.clone() |
| masks = F.interpolate(masks[None], size=(h, w))[0] |
|
|
| p_h, p_w = meta_data['padded_shape']['height'], meta_data['padded_shape']['width'] |
|
|
| padded_masks = torch.zeros(mask_cnt, p_h, p_w, dtype=masks.dtype) |
| padding = meta_data['padding'] |
|
|
| padded_masks[:, padding['before_height']:p_h - padding['after_height'], |
| padding['before_width']:p_w - padding['after_width']] = masks |
|
|
| |
| prompt_len = len(self.prompt) |
| labels = torch.ones_like(input_ids) * IGNORE_INDEX |
| labels[prompt_len:] = input_ids[prompt_len:] |
|
|
| if self.add_image_token: |
| input_ids[input_ids == self.image_token_idx] = IMAGE_TOKEN_INDEX |
|
|
| return dict(input_ids=input_ids, |
| mask_ids=mask_ids, |
| pixel_values=pixel_values, |
| padded_masks=padded_masks, |
| masks=masks, |
| gt_masks=gt_masks, |
| image_sizes=torch.tensor(image_data['image_sizes'][0]), |
| image=image, |
| meta_data=meta_data, |
| labels=labels) |
|
|
|
|
| if __name__ == '__main__': |
| from mmdet.datasets import RefCocoDataset |
| from mmengine.config import Config |
| from mmdet.datasets.transforms import LoadAnnotations |
|
|
| cfg = Config.fromfile('configs/fuyu/frozen_fuyu_8b_unet_sam_l_refcoco_png.py') |
| prompt_template = cfg.prompt_template |
| tokenizer = cfg.tokenizer |
| image_processor = cfg.image_processor |
| prompt = cfg.get('prompt', None) |
|
|
| refcoco2png_params = dict( |
| type=RefCOCO2PNG, |
| image_processor=image_processor, |
| tokenizer=tokenizer, |
| prompt_template=prompt_template, |
|
|
| ) |
| if prompt is not None: |
| refcoco2png_params.update(prompt=prompt) |
|
|
| test_pipeline = [ |
| dict(type=PILLoadImageFromFile, backend_args=None), |
| dict( |
| type=LoadAnnotations, |
| with_mask=True, |
| with_bbox=False, |
| with_seg=False, |
| with_label=False), |
| refcoco2png_params |
| ] |
|
|
| dataset = RefCocoDataset( |
| data_root='data/coco/', |
| data_prefix=dict(img_path='train2014/'), |
| text_mode='select_first', |
| pipeline=test_pipeline, |
| ann_file='refcoco/instances.json', |
| split_file='refcoco/refs(unc).p', |
| split='val' |
| ) |
|
|
|
|
| for data in dataset: |
| print(data.keys()) |
