Update app.py

app.py

@@ -1,83 +1,171 @@
-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()
+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):
+    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, 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()
+"""
+
+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()
+
+"""