image_prediction_ax.py

import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
# import onnxruntime as ort
import cv2
from utils.EdgeTAM_image_predictor import ImagePredictor
import argparse

np.random.seed(3)

def show_mask(mask, ax, random_color=False, borders = True):
    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 = mask.astype(np.uint8)
    mask_image =  mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    if borders:
        import cv2
        contours, _ = cv2.findContours(mask,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        # print(contours)
        # Try to smooth contours
        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
        # cv2.imwrite('./mask_image.jpg', mask_image)
    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_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))

def show_masks(
    image, 
    masks, 
    scores, 
    point_coords=None, 
    box_coords=None, 
    input_labels=None, 
    borders=True,
    save_dir="./results",        # 新增：保存目录
    base_name="mask"             # 新增：基础文件名
):
    """
    保存分割结果图像到文件，不再显示。
    
    Args:
        save_dir: 保存目录（会自动创建）
        base_name: 文件名前缀，如 "mask" → "mask_1.png"
    """
    os.makedirs(save_dir, exist_ok=True)
    
    for i, (mask, score) in enumerate(zip(masks, scores)):
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        show_mask(mask, plt.gca(), borders=borders)
        
        if point_coords is not None:
            assert input_labels is not None
            show_points(point_coords, input_labels, plt.gca())
            
        if box_coords is not None:
            show_box(box_coords, plt.gca())
            
        if len(scores) > 1:
            plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
            
        plt.axis('off')
        
        # 保存图像（不再 plt.show()）
        save_path = os.path.join(save_dir, f"{base_name}_{i+1}.png")
        plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=150)
        plt.close()  # 释放内存
        print(f"✅ Saved: {save_path}")


if __name__ == "__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument("--image_path", type=str, default="./examples/images/truck.jpg", help="Path to the input image.")
    argparser.add_argument("--model_path", type=str, default="./axmodel", help="Path to the ImagePredictor model.")
    argparser.add_argument("--save_dir", type=str, default="./results", help="Directory to save the output images.")
    argparser.add_argument("--input_box", type=str, default=None, help="Input box coordinates as x1,y1,x2,y2")
    argparser.add_argument("--input_mask", type=str, default=None, help="Path to the input mask numpy file.")
    argparser.add_argument("--input_point_coords", type=str, default=None, help="Input point coordinates as x1,y1 or x1,y1:x2,y2")
    argparser.add_argument("--input_point_labels", type=str, default=None, help="Input point labels as 1 or 0 or 1:0")

    args = argparser.parse_args()

    # load image
    image = np.array(Image.open(args.image_path).convert("RGB"))

    predictor = ImagePredictor(args.model_path)

    predictor.set_image(image)

    # define input prompts
    if args.input_mask is not None:
        input_mask = np.load(args.input_mask)
    else:
        input_mask = np.zeros((1, 256, 256), dtype=np.float32)

    if args.input_box is not None:
        input_box = np.array([int(x) for x in args.input_box.split(",")])
    else:
        input_box = None

    if args.input_point_coords is not None:
        input_point_coords = np.array([[int(coord) for coord in point.split(",")] for point in args.input_point_coords.split(":")])
    else:
        input_point_coords = None

    if args.input_point_labels is not None:
        input_point_labels = np.array([int(label) for label in args.input_point_labels.split(":")])
    else:
        input_point_labels = None

    if input_box is None and input_point_coords is None:
        raise ValueError("At least one of input_box or input_point_coords must be provided.")

    print("Get prompts: ")
    print(f"  input_box: {input_box}")
    print(f"  input_point_coords: {input_point_coords}")
    print(f"  input_point_labels: {input_point_labels}")

    #only box
    # input_box = np.array([75, 275, 1725, 850])
    # input_point_coords = None
    # input_point_labels = None

    # input_box = np.array([1375, 550, 1650, 800])
    # input_point_coords = None
    # input_point_labels = None

    #only point
    # input_box = None
    # input_point_coords = np.array([[500, 375], [1125, 625]])
    # input_point_labels = np.array([1, 1])

    # input_box = None
    # input_point_coords = np.array([[500, 375], [1125, 625]])
    # input_point_labels = np.array([1, 0])

    #point + box
    # input_box = np.array([425, 600, 700, 875])
    # input_point_coords = np.array([[575, 750]])
    # input_point_labels = np.array([0])

    # input_mask = np.load("./axmodel/logits.npy")

    # predict masks
    masks, scores, logits = predictor.predict(
        point_coords=input_point_coords,
        point_labels=input_point_labels,
        box=input_box,
        mask_input=input_mask,
        multimask_output=False,
    )

    sorted_ind = np.argsort(scores)[::-1]
    masks = masks[sorted_ind]
    scores = scores[sorted_ind]
    logits = logits[sorted_ind]
    print(scores)

    # visualize results
    show_masks(
        image,
        masks,
        scores,
        point_coords=input_point_coords,
        box_coords=input_box,
        input_labels=input_point_labels,
        borders=True,
    )