ymy-k / Hi-SAM

[arXiv preprint] Hi-SAM: Marrying Segment Anything Model for Hierarchical Text Segmentation
Apache License 2.0
193 stars 10 forks source link

How can I get a Paragraph Segmentation without giving a point prompt? #9

Closed greyzyl closed 3 months ago

greyzyl commented 4 months ago

such as ,this image Snipaste_2024-05-13_15-27-01

ymy-k commented 3 months ago

For example, when evaluating the model on HierText, a json file will be saved for each image. You can plot the results in the json file:

import json
from tqdm import tqdm
import random
import matplotlib.pyplot as plt
import cv2
import numpy as np

def show_mask(mask, ax, random_color=False, color=None):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.5])], axis=0)
    else:
        color = color if color is not None else np.array([30/255, 144/255, 255/255, 0.5])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

def draw_layout_results(img_id ='84f1fe431fcf0a31'):
    with open(img_id + '.jsonl') as f:  # modify your path
        img_res = json.load(f)
    image = cv2.imread(img_id + '.jpg')  # modify your path
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w = image.shape[:2]
    plt.figure(dpi=200)
    plt.imshow(image)
    print('start drawing')
    paras = img_res['paragraphs']
    for para in tqdm(paras):
        a = random.randint(1, 255)
        b = random.randint(1, 255)
        c = random.randint(1, 255)
        lines = para['lines']
        for line in lines:
            words = line['words']
            for word in words:
                poly = np.array(word['vertices'])
                mask = np.zeros((h, w), dtype=np.uint8)
                mask = cv2.fillPoly(mask, [poly], [1])
                show_mask(mask, plt.gca(), color=np.array([a / 255, b / 255, c / 255, 0.6]))
    plt.axis('off')
    plt.savefig(img_id+'_layout.png', bbox_inches='tight')
    plt.close()

draw_layout_results(your_img_id)