faster points

app.py

@@ -9,24 +9,18 @@ import matplotlib.pyplot as plt
 import torch
 from transformers import SamModel, SamProcessor
 
-import os
-
+from gradio_image_prompter import ImagePrompter
 
-# Define variables
-path = os.getcwd()
-font_path = r'{}/arial.ttf'.format(path)
-print(font_path)
+import os
 
-# Load the pre-trained model - FastSAM
-# fastsam_model = FastSAM('./FastSAM-s.pt')
+# define variables
 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
+# model_id = "facebook/sam-vit-huge" #60s
+model_id = 'Zigeng/SlimSAM-uniform-50' #50s
+# model_id = "facebook/sam-vit-base" #50s
+# model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+model = SamModel.from_pretrained(model_id).to(device)
+processor = SamProcessor.from_pretrained(model_id)
 
 # Description
 title = "<center><strong><font size='8'> 🍔 Segment food with clicks 🍜</font></strong></center>"
@@ -34,85 +28,23 @@ title = "<center><strong><font size='8'> 🍔 Segment food with clicks 🍜</fon
 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 2: Add positive (right click), negative (middle click), and bounding box (click and drag - only ONE box at most) 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%; }"
 
+# functions
 
 def read_image(url):
     response = requests.get(url)
     img = Image.open(BytesIO(response.content))
-    
-    global global_points
-    global global_point_label
-  
-    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
+    formatted_image = {
+        "image": np.array(img),
+        "points": [],
+    }  # Create the correct format
+    return formatted_image
 
 def get_mask_image(raw_image, mask):
     tmp_mask = np.array(mask * 1)
@@ -123,29 +55,32 @@ def get_mask_image(raw_image, mask):
     tmp_img_arr = np.concatenate((tmp_img_arr, tmp_mask2), axis = 2)
     return tmp_img_arr
 
+def format_prompt_points(points):
+    prompt_points = []
+    point_labels = []
+    prompt_boxes = []
+    for point in points:
+        print(point)
+        if point[2] == 2.0 and point[5] == 3.0:
+            prompt_boxes.append([point[0], point[1], point[3], point[4]])
+        else:
+            prompt_points.append([point[0], point[1]])
+            label = 1 if point[2] == 1.0 else 0
+            point_labels.append(label)
+    return prompt_points, point_labels, prompt_boxes
 
 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,
+    prompts
 ):
-    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)
+    image = np.array(prompts["image"])  # Convert the image to a numpy array
+    points = prompts["points"]  # Get the points from prompts
+    #
+    prompt_points, point_labels, prompt_boxes = format_prompt_points(points)
     print(prompt_points, point_labels, prompt_boxes)
     # segment
-    inputs = processor(raw_image, 
-                       input_boxes = prompt_boxes,
+    inputs = processor(image, 
+                       input_boxes = [prompt_boxes],
                        input_points=[[prompt_points]], 
                        input_labels=[point_labels], 
                        return_tensors="pt").to(device)
@@ -155,74 +90,15 @@ def segment_with_points(
     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_images = [get_mask_image(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)
-    
-    print(x, y, label)
-    print(previous_box_points)
-    
-    draw = ImageDraw.Draw(image)
-    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):
@@ -231,42 +107,32 @@ with gr.Blocks(css=css, title='Segment Food with Prompts') as demo:
             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")
+            segment_btn = gr.Button("Segment with prompts", variant='primary')
+            clear_btn = gr.Button("Clear points", variant='secondary')
         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()
+            candidate_pic = ImagePrompter(show_label=False)
+            segpic_output1 = gr.Image(format="png")
         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')
+            segpic_output2 = gr.Image(format="png")
+            segpic_output3 = gr.Image(format="png")
 
     # Define interaction relationship
     run_with_url.click(read_image,
                         inputs=[image_url],
                         # outputs=[segm_img_p, cond_img_p])
-                        outputs=[cond_img_p])
+                        outputs=[candidate_pic])
 
-    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],
+    segment_btn.click(segment_with_points,
+                        inputs=candidate_pic,
                         # outputs=[segm_img_p, cond_img_p])
-                        outputs=[segm_img_p1, segm_img_p2, segm_img_p3])
+                        outputs=[segpic_output1, segpic_output2, segpic_output3])
     
-    clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p1, segm_img_p2, segm_img_p3])
+    clear_btn.click(clear, outputs=[candidate_pic, segpic_output1, segpic_output2, segpic_output3])
 
 demo.queue()
 demo.launch()

app_orig_0215.py

@@ -0,0 +1,272 @@
+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_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)
+    
+    print(x, y, label)
+    print(previous_box_points)
+    
+    draw = ImageDraw.Draw(image)
+    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()