44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
import os
|
|
from ultralytics import YOLO
|
|
|
|
# 模型缓存
|
|
_yolo_models = {}
|
|
|
|
|
|
def detect_crop_area(image_path: str, model_path: str):
|
|
"""
|
|
使用YOLOv8模型检测图像中的裁切区域。
|
|
|
|
Args:
|
|
image_path (str): 原始图像的文件路径。
|
|
model_path (str): 用于检测的YOLOv8模型 (.pt) 的路径。
|
|
|
|
Returns:
|
|
tuple or None: 如果检测到物体，返回 (x1, y1, x2, y2)；否则返回 None。
|
|
"""
|
|
if model_path not in _yolo_models:
|
|
print(f"Loading YOLOv8 model from: {model_path}")
|
|
if not os.path.exists(model_path):
|
|
print(f"Error: YOLO model file not found at {model_path}")
|
|
return None
|
|
_yolo_models[model_path] = YOLO(model_path)
|
|
|
|
model = _yolo_models[model_path]
|
|
|
|
if not os.path.exists(image_path):
|
|
print(f"Error: Image file not found at {image_path}")
|
|
return None
|
|
|
|
try:
|
|
results = model.predict(source=image_path, conf=0.5, verbose=False)
|
|
|
|
if not results or not results[0].boxes:
|
|
print(f"Warning: YOLO did not detect any objects in {image_path}")
|
|
return None
|
|
|
|
box = results[0].boxes.xyxy[0].cpu().numpy().astype(int)
|
|
return tuple(box)
|
|
|
|
except Exception as e:
|
|
print(f"An error occurred during YOLO prediction for {image_path}: {e}")
|
|
return None |