Method parse_dense_caption

caption_anything/model.py:219–268

(self, image, topN=10, reference_caption=[], verbose=False)

        return out_list

    def parse_dense_caption(self, image, topN=10, reference_caption=[], verbose=False):
        width, height = get_image_shape(image)
        prompt = {'prompt_type': ['everything']}
        densecap_args = {
            'return_ppl': True,
            'clip_filter': True,
            'reference_caption': reference_caption,
            'text_prompt': "",  # 'Question: what does the image show? Answer:'
            'seg_crop_mode': 'w_bg',
            # 'text_prompt': "",
            # 'seg_crop_mode': 'wo_bg',
            'disable_regular_box': False,
            'topN': topN,
            'min_ppl_score': -1.8,
            'min_clip_score': 0.30,
            'min_mask_area': 2500,
        }

        dense_captions = self.inference(image, prompt,
                                        controls=None,
                                        disable_gpt=True,
                                        verbose=verbose,
                                        is_densecap=True,
                                        args=densecap_args)
        print('Process Dense Captioning: \n', dense_captions)
        dense_captions = list(filter(
            lambda x: x['ppl_score'] / (1 + len(x['generated_captions']['raw_caption'].split())) >= densecap_args[
                'min_ppl_score'], dense_captions))
        dense_captions = list(filter(lambda x: x['clip_score'] >= densecap_args['min_clip_score'], dense_captions))
        dense_cap_prompt = []
        for cap in dense_captions:
            x, y, w, h = cap['bbox']
            cx, cy = x + w / 2, (y + h / 2)
            dense_cap_prompt.append(
                "({}: X:{:.0f}, Y:{:.0f}, Width:{:.0f}, Height:{:.0f})".format(cap['generated_captions']['raw_caption'],
                                                                               cx, cy, w, h))

        if verbose:
            all_masks = [np.array(item['mask'].convert('P')) for item in dense_captions]
            new_image = mask_painter_foreground_all(np.array(image), all_masks, background_alpha=0.4)
            save_path = 'result/dense_caption_mask.png'
            Image.fromarray(new_image).save(save_path)
            print(f'Dense captioning mask saved in {save_path}')

            vis_path = 'result/dense_caption_vis_{}.png'.format(time.time())
            dense_cap_painter_input = [{'bbox': xywh_to_x1y1x2y2(cap['bbox']),
                                        'caption': cap['generated_captions']['raw_caption']} for cap in dense_captions]
            draw_bbox(load_image(image, return_type='numpy'), vis_path, dense_cap_painter_input, show_caption=True)
            print(f'Dense Captioning visualization saved in {vis_path}')
        return ','.join(dense_cap_prompt)

    def parse_ocr(self, image, thres=0.2):
        width, height = get_image_shape(image)