import os
import time
import numpy as np
import torch
from tqdm import tqdm
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
from PIL import Image

assert torch.cuda.is_available()
device = torch.device("cuda")

torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
if torch.cuda.get_device_properties(0).major >= 8:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

sam2_checkpoint = "/workspace/sam2/checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"

sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device, apply_postprocessing=False)

mask_generator = SAM2AutomaticMaskGenerator(
    model=sam2_model,
    #img_resize_size=640,
    # points_per_side=64,
    # points_per_batch=128,
    # pred_iou_thresh=0.7,
    # stability_score_thresh=0.92,
    # stability_score_offset=0.7,
    # crop_n_layers=1,
    # box_nms_thresh=0.7,
    # crop_n_points_downscale_factor=2,
    # min_mask_region_area=25.0,
    # use_m2m=True,
)

# 图像文件夹路径
images_dir = "/workspace/sam2/coco_mini/images/val2017"

images_names = [
    p for p in os.listdir(images_dir)
    if os.path.splitext(p)[-1].lower() in [".jpg", ".jpeg"]
]
images_names.sort(key=lambda p: int(os.path.splitext(p)[0]))

images = []
for frame_name in images_names:
    img_path = os.path.join(images_dir, frame_name)
    img = np.array(Image.open(img_path).convert("RGB"))
    images.append(img)

num_frames = len(images)
print(f"load {num_frames} images from {images_dir}")


torch.cuda.empty_cache()

print("warm up...")
with torch.inference_mode():
    with torch.autocast("cuda", dtype=torch.bfloat16):
        masks = mask_generator.generate(images[0])
        masks = mask_generator.generate(images[1])

print("process images...")
total_start = time.time()

mask_counts = []

with torch.inference_mode():
    with torch.autocast("cuda", dtype=torch.bfloat16):
        for img in tqdm(images, desc="process images"):
            masks = mask_generator.generate(img)
            mask_counts.append(len(masks))

total_end = time.time()
total_time = total_end - total_start

# 计算结果
avg_time_per_sample = total_time / num_frames
avg_masks_per_sample = np.mean(mask_counts) if mask_counts else 0

print("\n=== Result ===")
print(f"Total time: {total_time:.2f} s")
print(f"Iference images: {num_frames}")
print(f"Average time per sample: {avg_time_per_sample:.2f} s/sample")
