Update app.py

app.py

@@ -1,22 +1,9 @@
 import gradio as gr
-import sahi
 import torch
-from ultralyticsplus import *
-
-# Images
-sahi.utils.file.download_from_url(
-    "https://raw.githubusercontent.com/kadirnar/dethub/main/data/images/highway.jpg",
-    "highway.jpg",
-)
-sahi.utils.file.download_from_url(
-    "https://raw.githubusercontent.com/obss/sahi/main/tests/data/small-vehicles1.jpeg",
-    "small-vehicles1.jpeg",
-)
-sahi.utils.file.download_from_url(
-    "https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/zidane.jpg",
-    "zidane.jpg",
-)
-
+from ultralyticsplus import YOLO, render_result
+import numpy as np
+from PIL import Image
+import cv2
 
 model_names = [
     "yolov8n-seg.pt",
@@ -29,13 +16,70 @@ model_names = [
 current_model_name = "yolov8m-seg.pt"
 model = YOLO(current_model_name)
 
+def sort_instance_masks_by_centroid(instances_mask, reverse=False):
+    # Calculate centroid of each instance mask
+    centroids = []
+    for mask in instances_mask:
+        # Find contours of the mask
+        mask_np = mask.astype(np.uint8)
+        #mask_np[mask_np !=0] = 255
+        contours, hierarchy = cv2.findContours(mask_np, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+        # Calculate moments of the contour
+        moments = cv2.moments(contours[0])
+
+        # Calculate centroid coordinates
+        c_x = int(moments["m10"] / moments["m00"])
+        c_y = int(moments["m01"] / moments["m00"])
+        centroids.append((c_x, c_y))
+
+    # Sort instance masks by centroid coordinates
+    sorted_instances_mask = [instance_mask for _, instance_mask in
+                                     sorted(zip(centroids, instances_mask), reverse=reverse)]
+
+    return sorted_instances_mask
+
+def visualize_masks(masks):
+    masks = masks.detach().cpu().numpy()
+    height, width = masks.shape[1:]
+    # 计算有多少个 mask
+    num_masks = masks.shape[0]
+    masks = sort_instance_masks_by_centroid(masks)
+
+
+    # 创建一个空白图像，背景颜色为黑色
+    img = Image.new('RGB', (width, height),(0,0,0))
+    #img.putpalette([0, 0, 0] * 256)
+    img_array = np.array(img)
+    colors = []
+
+    # 将每个 mask 标记为不同的颜色
+    for i in range(num_masks):
+        color = np.random.randint(0, 256, size=3)
+        colors.append(tuple(color))
+        #colorimg.paste
+        #colorimg = Image.new('RGB', (width,height), color=tuple(np.random.randint(0, 256, size=3)))
+        #mask_img_tmp = Image.fromarray(masks[i]).convert('RGB')
+        #mask_array = Image.fromarray(masks[i])
+        img_array[masks[i] != 0,:] = color
+        #mask_img = mask_img.putpalette(color)
+        #img.paste(mask_img,(0,0),mask_img_tmp)
+
+        #img.putpalette(color + (0,) * 253)
+
+    # 将 mask 根据颜色映射显示为 RGB 图像
+    img_rgb = Image.fromarray(img_array)
+    return img_rgb,colors
+
+
 
 def yolov8_inference(
-    image: gr.inputs.Image = None,
-    model_name: gr.inputs.Dropdown = None,
-    image_size: gr.inputs.Slider = 640,
-    conf_threshold: gr.inputs.Slider = 0.25,
-    iou_threshold: gr.inputs.Slider = 0.45,
+    image = None,
+    model_name = None,
+    dest_width = 512,
+    dest_height = 512,
+    conf_threshold = 0.25,
+    iou_threshold = 0.45,
 ):
     """
     YOLOv8 inference function
@@ -55,16 +99,21 @@ def yolov8_inference(
         current_model_name = model_name
     model.overrides["conf"] = conf_threshold
     model.overrides["iou"] = iou_threshold
-    results = model.predict(image, imgsz=image_size, return_outputs=True)
+    model.overrides["classes"] = [0]
+    results = model.predict(image)
     renders = []
-    for image_results in model.predict(image, imgsz=image_size, return_outputs=True):
-        render = render_model_output(
-            model=model, image=image, model_output=image_results
-        )
+    colorarray = []
+    for image_results in model.predict(image):
+        #print("predict results:  ",type(image_results.masks))
+        #render = render_result(
+        #    model=model, image=image, result=image_results
+        #)
+        render ,colors= visualize_masks(image_results.masks.data)
+        render = render.resize((dest_width,dest_height))
         renders.append(render)
+        colorarray.append(colors)
 
-    return renders[0]
-
+    return renders[0],','.join(['#%02x%02x%02x' % row for row in colorarray[0]])
 
 inputs = [
     gr.Image(type="filepath", label="Input Image"),
@@ -73,28 +122,26 @@ inputs = [
         value=current_model_name,
         label="Model type",
     ),
-    gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
+    gr.inputs.Slider(minimum=128, maximum=2048, step=64, default=512, label="Width"),
+    gr.inputs.Slider(minimum=128, maximum=2048, step=64, default=512, label="Height"),
+
     gr.Slider(
         minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"
     ),
     gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
 ]
 
-outputs = gr.Image(type="filepath", label="Output Image")
-title = "Ultralytics YOLOv8 Segmentation Demo"
+outputs = [gr.Image(type="filepath", label="Output Image"),gr.Textbox(label="Output Text")]
+title = "Ultralytics YOLOv8 Segmentation For HumanBody Only Now"
+
 
-examples = [
-    ["zidane.jpg", "yolov8m-seg.pt", 640, 0.6, 0.45],
-    ["highway.jpg", "yolov8m-seg.pt", 640, 0.25, 0.45],
-    ["small-vehicles1.jpeg", "yolov8m-seg.pt", 640, 0.25, 0.45],
-]
 demo_app = gr.Interface(
     fn=yolov8_inference,
     inputs=inputs,
     outputs=outputs,
     title=title,
-    examples=examples,
-    cache_examples=True,
+    examples=None,
+    cache_examples=False,
     theme="default",
 )
-demo_app.launch(debug=True, enable_queue=True)
+demo_app.launch(debug=True, enable_queue=True)