SAHI with YOLOv5 for Sliced Inference

inference_for_yolov5.ipynb - Colab
bash 复制代码
# arrange an instance segmentation model for test
import torch

# 保存原始函数
_original_torch_load = torch.load

def _patched_torch_load(*args, **kwargs):
    # 强制设置 weights_only=False
    kwargs['weights_only'] = False
    return _original_torch_load(*args, **kwargs)

# 替换 torch.load
torch.load = _patched_torch_load

from IPython.display import Image

# import required functions, classes
from sahi import AutoDetectionModel
from sahi.predict import get_prediction, get_sliced_prediction, predict
from sahi.utils.cv import read_image
from sahi.utils.file import download_from_url
from sahi.utils.yolov5 import download_yolov5s6_model

# download YOLOV5S6 model to 'models/yolov5s6.pt'
yolov5_model_path = "models/yolov5s6.pt"
download_yolov5s6_model(destination_path=yolov5_model_path)

# download test images into demo_data folder
# download_from_url(
#     "https://raw.githubusercontent.com/obss/sahi/main/demo/demo_data/small-vehicles1.jpeg",
#     "demo_data/small-vehicles1.jpeg",
# )
# download_from_url(
#     "https://raw.githubusercontent.com/obss/sahi/main/demo/demo_data/terrain2.png", "demo_data/terrain2.png"
# )

detection_model = AutoDetectionModel.from_pretrained(
    model_type="yolov5",
    model_path=yolov5_model_path,
    confidence_threshold=0.3,
    device="cuda:0",  # or 'cuda:0'
)

# result = get_prediction("demo_data/small-vehicles1.jpeg", detection_model)

# result.export_visuals(export_dir="demo_data/")

# Image("demo_data/prediction_visual.png")


result = get_sliced_prediction(
    "demo_data/small-vehicles1.jpeg",
    detection_model,
    slice_height=256,
    slice_width=256,
    overlap_height_ratio=0.2,
    overlap_width_ratio=0.2,
)

result.export_visuals(export_dir="demo_data/")

Image("demo_data/prediction_sliced.png")


object_prediction_list = result.object_prediction_list
object_prediction_list[0]


print(result.to_coco_annotations()[:3])
print(result.to_coco_predictions(image_id=1)[:3])
print(result.to_imantics_annotations()[:3])
print(result.to_fiftyone_detections()[:3])


model_type = "yolov5"
model_path = yolov5_model_path
model_device = "cuda:0"  # or 'cpu'
model_confidence_threshold = 0.4

slice_height = 256
slice_width = 256
overlap_height_ratio = 0.2
overlap_width_ratio = 0.2

source_image_dir = "demo_data/"

predict(
    model_type=model_type,
    model_path=model_path,
    model_device=model_device,
    model_confidence_threshold=model_confidence_threshold,
    source=source_image_dir,
    slice_height=slice_height,
    slice_width=slice_width,
    overlap_height_ratio=overlap_height_ratio,
    overlap_width_ratio=overlap_width_ratio,
)