app.py

from ultralytics import YOLO
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
import cv2
import torch

model = YOLO('checkpoints/FastSAM.pt')  # load a custom model


def fast_process(annotations, image):
    fig = plt.figure(figsize=(10, 10))
    plt.imshow(image)
    #original_h = image.shape[0]
    #original_w = image.shape[1]
    #for i, mask in enumerate(annotations):
    #        mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
    #        annotations[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
    fast_show_mask(annotations,
                       plt.gca())
                       #target_height=original_h,
                       #target_width=original_w)
    plt.axis('off')
    plt.tight_layout()
    return fig


#   CPU post process
def fast_show_mask(annotation, ax):
    msak_sum = annotation.shape[0]
    height = annotation.shape[1]
    weight = annotation.shape[2]
    # 将annotation 按照面积 排序
    areas = np.sum(annotation, axis=(1, 2))
    sorted_indices = np.argsort(areas)[::1]
    annotation = annotation[sorted_indices]

    index = (annotation != 0).argmax(axis=0)
    color = np.random.random((msak_sum, 1, 1, 3))
    transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
    visual = np.concatenate([color, transparency], axis=-1)
    mask_image = np.expand_dims(annotation, -1) * visual

    show = np.zeros((height, weight, 4))

    h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij')
    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
    # 使用向量化索引更新show的值
    show[h_indices, w_indices, :] = mask_image[indices]


    #if retinamask == False:
    #    show = cv2.resize(show, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
    ax.imshow(show)



# post_process(results[0].masks, Image.open("../data/cake.png"))

def predict(input, input_size=512):
    input_size = int(input_size)  # 确保 imgsz 是整数
    results = model(input, device='cpu', retina_masks=True, iou=0.7, conf=0.25, imgsz=input_size)
    pil_image = fast_process(annotations=results[0].masks.data.numpy(), image=input)

    return pil_image


# inp = 'assets/sa_192.jpg'
# results = model(inp, device='cpu', retina_masks=True, iou=0.7, conf=0.25, imgsz=1024)
# results = format_results(results[0], 100)
# post_process(annotations=results, image_path=inp)

demo = gr.Interface(fn=predict,
                    inputs=[gr.inputs.Image(type='pil'), gr.inputs.Dropdown(choices=[512, 800, 1024], default=512)],
                    outputs=['plot'],
                    examples=[["assets/sa_8776.jpg", 1024]],
                    #    ["assets/sa_1309.jpg", 1024]],
                    # examples=[["assets/sa_192.jpg"], ["assets/sa_414.jpg"],
                    #           ["assets/sa_561.jpg"], ["assets/sa_862.jpg"],
                    #           ["assets/sa_1309.jpg"], ["assets/sa_8776.jpg"],
                    #           ["assets/sa_10039.jpg"], ["assets/sa_11025.jpg"],],
                    )

demo.launch()
"""



from ultralytics import YOLO

import numpy as np

import matplotlib.pyplot as plt

import gradio as gr

import torch



model = YOLO('checkpoints/FastSAM.pt')  # load a custom model



def format_results(result,filter = 0):

    annotations = []

    n = len(result.masks.data)

    for i in range(n):

        annotation = {}

        mask = result.masks.data[i] == 1.0



        if torch.sum(mask) < filter:

            continue

        annotation['id'] = i

        annotation['segmentation'] = mask.cpu().numpy()

        annotation['bbox'] = result.boxes.data[i]

        annotation['score'] = result.boxes.conf[i]

        annotation['area'] = annotation['segmentation'].sum()

        annotations.append(annotation)

    return annotations



def show_mask(annotation, ax, random_color=True, bbox=None, points=None):

    if random_color :    # random mask color

        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)

    else:

        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])

    if type(annotation) == dict:

        annotation = annotation['segmentation']

    mask = annotation

    h, w = mask.shape[-2:]

    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)

    # draw box

    if bbox is not None:

        x1, y1, x2, y2 = bbox

        ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))

    # draw point

    if points is not None:

        ax.scatter([point[0] for point in points], [point[1] for point in points], s=10, c='g')

    ax.imshow(mask_image)

    return mask_image



def post_process(annotations, image, mask_random_color=True, bbox=None, points=None):

    fig = plt.figure(figsize=(10, 10))

    plt.imshow(image)

    for i, mask in enumerate(annotations):

        show_mask(mask, plt.gca(),random_color=mask_random_color,bbox=bbox,points=points)

    plt.axis('off')

    

    plt.tight_layout()

    return fig





# post_process(results[0].masks, Image.open("../data/cake.png"))



def predict(input, input_size):

    input_size = int(input_size)  # 确保 imgsz 是整数

    results = model(input, device='cpu', retina_masks=True, iou=0.7, conf=0.25, imgsz=input_size)

    results = format_results(results[0], 100)

    results.sort(key=lambda x: x['area'], reverse=True)

    pil_image = post_process(annotations=results, image=input)

    return pil_image



# inp = 'assets/sa_192.jpg'

# results = model(inp, device='cpu', retina_masks=True, iou=0.7, conf=0.25, imgsz=1024)

# results = format_results(results[0], 100)

# post_process(annotations=results, image_path=inp)



demo = gr.Interface(fn=predict,

                    inputs=[gr.inputs.Image(type='pil'), gr.inputs.Dropdown(choices=[512, 800, 1024], default=1024)],

                    outputs=['plot'],

                     examples=[["assets/sa_8776.jpg", 1024]],

                            #    ["assets/sa_1309.jpg", 1024]],

                    # examples=[["assets/sa_192.jpg"], ["assets/sa_414.jpg"],

                    #           ["assets/sa_561.jpg"], ["assets/sa_862.jpg"],

                    #           ["assets/sa_1309.jpg"], ["assets/sa_8776.jpg"],

                    #           ["assets/sa_10039.jpg"], ["assets/sa_11025.jpg"],],

                    )



demo.launch()



"""