Update app.py

app.py

@@ -38,33 +38,21 @@ def fn(image):
         pred_score_thr=instance_thres
     )
 
-    drawed = img.copy()
-    im_h, im_w = img.shape[:2]
+    # 创建一个空白的白色图像，和原图大小一致
+    white = np.full_like(img, 255)
 
     # instances.bboxes, instances.masks will be None, None if no obj is detected
     if instances.bboxes is None:
-        return Image.fromarray(drawed[..., ::-1])
+        return Image.fromarray(white)
 
     for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
-        color = get_color(ii)
+        # 把mask转换为bool类型，方便后续操作
+        mask = mask.astype(np.bool)
 
-        mask_alpha = 0.5
-        linewidth = max(round(sum(img.shape) / 2 * 0.003), 2)
+        # 用原图中对应的区域替换白色图像中的区域，实现去除背景的效果
+        white[mask] = img[mask]
 
-        # draw bbox
-        p1, p2 = (int(xywh[0]), int(xywh[1])), (int(xywh[2] + xywh[0]), int(xywh[3] + xywh[1]))
-        cv2.rectangle(drawed, p1, p2, color, thickness=linewidth, lineType=cv2.LINE_AA)
-        
-        # draw mask
-        p = mask.astype(np.float32)
-        blend_mask = np.full((im_h, im_w, 3), color, dtype=np.float32)
-        alpha_msk = (mask_alpha * p)[..., None]
-        alpha_ori = 1 - alpha_msk
-        drawed = drawed * alpha_ori + alpha_msk * blend_mask
-
-        drawed = drawed.astype(np.uint8)
-
-    return Image.fromarray(drawed[..., ::-1])
+    return Image.fromarray(white[..., ::-1])
 
 iface = gr.Interface(
     # design titles and text descriptions