app.py

import gradio as gr
import torch
from PIL import ImageDraw, Image, ImageFont
import numpy as np
import requests
from io import BytesIO

import matplotlib.pyplot as plt
import torch
from transformers import SamModel, SamProcessor

import os


# Define variables
path = os.getcwd()
font_path = r'{}/arial.ttf'.format(path)
print(font_path)

# Load the pre-trained model - FastSAM
# fastsam_model = FastSAM('./FastSAM-s.pt')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")

# Points
global_points = []
global_point_label = []
previous_box_points = 0

# Description
title = "<center><strong><font size='8'> 🍔 Segment food with clicks 🍜</font></strong></center>"

instruction = """ # Instruction
        This segmentation tool is built with HuggingFace SAM model. To use to label true mask, please follow the following steps \n
        🔥 Step 1: Copy segmentation candidate image link and paste in 'Enter Image URL' and click 'Upload Image' \n
        🔥 Step 2: Add positive (Add mask), negative (Remove Area), and bounding box for the food \n
        🔥 Step 3: Click on 'Segment with prompts' to segment Image and see if there's a correct segmentation on the 3 options \n
        🔥 Step 4: If not, you can repeat the process of adding prompt and segment until a correct one is generated. Prompt history will be retained unless reloading the image \n
        🔥 Step 5: Download the satisfied segmentaion image through the icon on top right corner of the image, please name it with 'correct_seg_xxx' where xxx is the photo ID 
        """  
css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"


def read_image(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    
    global global_points
    global global_point_label
    global draw
  
    global_points = []
    global_point_label = []

    return img

# def show_mask(mask, ax, random_color=False):
#     if random_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])
#     h, w = mask.shape[-2:]
#     mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
#     ax.imshow(mask_image)

# def show_points(coords, labels, ax, marker_size=375):
#     pos_points = coords[labels==1]
#     neg_points = coords[labels==0]
#     ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
#     ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)

# def show_masks_and_points_on_image(raw_image, mask, input_points, input_labels, args):
#     masks = masks.squeeze() if len(masks.shape) == 4 else masks.unsqueeze(0) if len(masks.shape) == 2 else masks
#     scores = scores.squeeze() if (scores.shape[0] == 1) & (len(scores.shape) == 3) else scores
#     #
#     input_points = np.array(input_points)
#     labels = np.array(input_labels)
#     # 
#     mask = mask.cpu().detach()
#     plt.imshow(np.array(raw_image))
#     ax = plt.gca()
#     show_mask(mask, ax)
#     show_points(input_points, labels, ax, marker_size=375)  
#     ax.axis("off")    

#     save_path = args.output
#     if not os.path.exists(save_path):
#         os.makedirs(save_path)
#     plt.axis("off")
#     fig = plt.gcf()
#     plt.draw()
    
#     try:
#         buf = fig.canvas.tostring_rgb()
#     except AttributeError:
#         fig.canvas.draw()
#         buf = fig.canvas.tostring_rgb()
    
#     cols, rows = fig.canvas.get_width_height()
#     img_array = np.fromstring(buf, dtype=np.uint8).reshape(rows, cols, 3)
#     cv2.imwrite(os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))

def format_prompt_points(points, labels):
    prompt_points = [xy for xy, l in zip(points, labels) if l != 9]
    point_labels = [l for l in labels if l != 9]
    #
    prompt_boxes = None
    if len(point_labels) < len(labels):
      prompt_boxes = [[np.array([xy for xy, l in zip(points, labels) if l == 9]).reshape(-1, 4).tolist()]]
    return prompt_points, point_labels, prompt_boxes

# def get_mask_image(raw_image, mask):
#     tmp_mask = np.array(mask)
#     tmp_img_arr = np.array(raw_image)
#     tmp_img_arr[tmp_mask == False] = [255,255,255]
#     return tmp_img_arr

def get_mask_image(raw_image, mask):
    tmp_mask = np.array(mask * 1)
    tmp_mask[tmp_mask == 1] = 255
    tmp_mask2 = np.expand_dims(tmp_mask, axis=2)
    #
    tmp_img_arr = np.array(raw_image)
    tmp_img_arr = np.concatenate((tmp_img_arr, tmp_mask2), axis = 2)
    return tmp_img_arr


def segment_with_points(
    input,
    input_size=1024, 
    iou_threshold=0.7,
    conf_threshold=0.25,
    better_quality=False,
    withContours=True,
    use_retina=True,
    mask_random_color=True,
):
    global global_points
    global global_point_label

    # read image
    raw_image = Image.open(requests.get(input, stream=True).raw).convert("RGB")

    # get prompts
    prompt_points, point_labels, prompt_boxes = format_prompt_points(global_points, global_point_label)
    print(prompt_points, point_labels, prompt_boxes)
    # segment
    inputs = processor(raw_image, 
                       input_boxes = prompt_boxes,
                       input_points=[[prompt_points]], 
                       input_labels=[point_labels], 
                       return_tensors="pt").to(device)
    with torch.no_grad():
      outputs = model(**inputs)
    #
    masks = processor.image_processor.post_process_masks(
        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
    scores = outputs.iou_scores

    # only show the first mask
    # fig = show_masks_and_points_on_image(raw_image, masks[0][0][0], [global_points], global_point_label)
    mask_images = [get_mask_image(raw_image, m) for m in masks[0][0]]
    mask_img1, mask_img2, mask_img3 = mask_images
    # return fig, None
    return mask_img1, mask_img2, mask_img3

def find_font_size(text, font_path, image, target_width_ratio):
    tested_font_size = 100
    tested_font = ImageFont.truetype(font_path, tested_font_size)
    observed_width = get_text_size(text, image, tested_font)
    estimated_font_size = tested_font_size / (observed_width / image.width) * target_width_ratio
    return round(estimated_font_size)

def get_text_size(text, image, font):
    im = Image.new('RGB', (image.width, image.height))
    draw = ImageDraw.Draw(im)
    return draw.textlength(text, font)


def get_points_with_draw(image, label, evt: gr.SelectData):
    global global_points
    global global_point_label
    global previous_box_points

    x, y = evt.index[0], evt.index[1]
    point_radius = 15
    point_color = (255, 255, 0) if label == 'Add Mask' else (255, 0, 255)
    global_points.append([x, y])
    global_point_label.append(1 if label == 'Add Mask' else 0 if label == 'Remove Area' else 9)

    draw = ImageDraw.Draw(image)
    
    print(x, y, label)
    print(previous_box_points)
    
    if label != 'Bounding Box':
        draw.ellipse([(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], fill=point_color)
    else:
        if (previous_box_points == 0) | (previous_box_points%2 == 0):
          target_width_ratio = 0.9
          text = "Please Click Another Point For Bounding Box"
          font_size = find_font_size(text, font_path, image, target_width_ratio)
          font = ImageFont.truetype(font_path, font_size)
          draw.text((x, y), text, fill = (0,0,0), font = font)
        else:
          [previous_x, previous_y] = global_points[-2]
          print((previous_x, previous_y), (x, y))
          draw.rectangle([(previous_x, previous_y), (x, y)], outline=(0, 0, 255), width=10)
        previous_box_points += 1
    return image

def clear():
    global global_points
    global global_point_label

    global_points = []
    global_point_label = []
    previous_box_points = 0
    return None, None, None, None


# Configure layout
cond_img_p = gr.Image(label="Input with points", type='pil')
segm_img_p1 = gr.Image(label="Segmented Image Option 1", interactive=False, type='pil', format="png")
segm_img_p2 = gr.Image(label="Segmented Image Option 2", interactive=False, type='pil', format="png")
segm_img_p3 = gr.Image(label="Segmented Image Option 3", interactive=False, type='pil', format="png")

with gr.Blocks(css=css, title='Segment Food with Prompts') as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(title)
            gr.Markdown("")
            image_url = gr.Textbox(label="Enter Image URL", 
                value = "https://img.cdn4dd.com/u/media/4da0fbcf-5e3d-45d4-8995-663fbcf3c3c8.jpg")
            run_with_url = gr.Button("Upload Image")
        with gr.Column(scale=1):
            gr.Markdown(instruction)

        # Images
    with gr.Row(variant="panel"):
        with gr.Column(scale=0):
            cond_img_p.render()
            segm_img_p2.render()
        with gr.Column(scale=0):
            segm_img_p1.render()
            segm_img_p3.render()
            
    # Submit & Clear
    with gr.Row():
        with gr.Column():
            add_or_remove = gr.Radio(["Add Mask", "Remove Area", "Bounding Box"], 
                    value="Add Mask", 
                    label="Point label")
        with gr.Column():
            segment_btn_p = gr.Button("Segment with prompts", variant='primary')
            clear_btn_p = gr.Button("Clear points", variant='secondary')

    # Define interaction relationship
    run_with_url.click(read_image,
                        inputs=[image_url],
                        # outputs=[segm_img_p, cond_img_p])
                        outputs=[cond_img_p])

    cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove], cond_img_p)

    segment_btn_p.click(segment_with_points,
                        inputs=[image_url],
                        # outputs=[segm_img_p, cond_img_p])
                        outputs=[segm_img_p1, segm_img_p2, segm_img_p3])
    
    clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p1, segm_img_p2, segm_img_p3])

demo.queue()
demo.launch()