70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
|
# SAM segmentation module for fish detection
|
|||
|
|
# Uses SAM (Segment Anything Model) for segmentation with boxes from YOLO detection
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import os
|
|||
|
|
import argparse
|
|||
|
|
import urllib.request
|
|||
|
|
from segment_anything import sam_model_registry, SamPredictor
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
|
|||
|
|
def download_models():
|
|||
|
|
"""
|
|||
|
|
下载所需的模型文件
|
|||
|
|
"""
|
|||
|
|
# 下载 SAM 模型
|
|||
|
|
sam_model_urls = {
|
|||
|
|
"vit_h": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
|
|||
|
|
"vit_l": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
|
|||
|
|
"vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 下载 SAM 模型
|
|||
|
|
sam_checkpoint = "sam_vit_h_4b8939.pth"
|
|||
|
|
if not os.path.exists(sam_checkpoint):
|
|||
|
|
print(f"正在下载SAM模型...")
|
|||
|
|
urllib.request.urlretrieve(sam_model_urls["vit_h"], sam_checkpoint)
|
|||
|
|
print(f"SAM模型下载完成")
|
|||
|
|
|
|||
|
|
def init_models(device="cpu", seg_model="sam", yolo_seg_weights=None):
|
|||
|
|
"""
|
|||
|
|
初始化模型
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
device: 运行设备
|
|||
|
|
seg_model: 分割模型类型 ("sam" 或 "yolov8_seg")
|
|||
|
|
yolo_seg_weights: YOLOv8分割模型权重路径（当seg_model="yolov8_seg"时必需）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
sam_predictor 或 yolo_seg_model: 分割模型
|
|||
|
|
"""
|
|||
|
|
if seg_model == "yolov8_seg":
|
|||
|
|
# 初始化YOLOv8分割模型
|
|||
|
|
if yolo_seg_weights is None:
|
|||
|
|
raise ValueError("使用yolov8_seg时必须提供yolo_seg_weights参数")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
except ImportError:
|
|||
|
|
raise ImportError("未找到 ultralytics，请先安装: pip install ultralytics")
|
|||
|
|
|
|||
|
|
print(f"正在加载YOLOv8分割模型: {yolo_seg_weights}")
|
|||
|
|
yolo_seg_model = YOLO(yolo_seg_weights)
|
|||
|
|
print("YOLOv8分割模型加载成功")
|
|||
|
|
return yolo_seg_model
|
|||
|
|
else:
|
|||
|
|
# 默认使用SAM
|
|||
|
|
# 下载模型
|
|||
|
|
download_models()
|
|||
|
|
|
|||
|
|
# 初始化SAM
|
|||
|
|
print("正在加载SAM模型...")
|
|||
|
|
sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
|
|||
|
|
sam.to(device=device)
|
|||
|
|
sam_predictor = SamPredictor(sam)
|
|||
|
|
print("SAM模型加载成功")
|
|||
|
|
return sam_predictor
|