181 lines
5.9 KiB
Python
181 lines
5.9 KiB
Python
import os
|
||
import tempfile
|
||
from typing import List
|
||
from urllib.parse import urlparse
|
||
|
||
import cv2
|
||
import requests
|
||
from pydantic import BaseModel
|
||
from tqdm import tqdm
|
||
from ultralytics import YOLO
|
||
|
||
|
||
# 定义返回值模型
|
||
class DetectionResponse(BaseModel):
|
||
hasTarget: int
|
||
originalImgSize: List[int]
|
||
targets: List[dict]
|
||
processing_errors: List[str] = []
|
||
|
||
|
||
def download_large_file(url, chunk_size=1024 * 1024):
|
||
"""下载大型文件到临时文件、返回临时文件路径"""
|
||
try:
|
||
response = requests.get(url, stream=True, timeout=30)
|
||
response.raise_for_status()
|
||
|
||
file_size = int(response.headers.get('Content-Length', 0))
|
||
|
||
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png') # 适配png格式
|
||
temp_file_path = temp_file.name
|
||
temp_file.close()
|
||
|
||
with open(temp_file_path, 'wb') as f, tqdm(
|
||
total=file_size, unit='B', unit_scale=True,
|
||
desc=f"下载 {os.path.basename(urlparse(url).path)}"
|
||
) as pbar:
|
||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||
if chunk:
|
||
f.write(chunk)
|
||
pbar.update(len(chunk))
|
||
|
||
return temp_file_path
|
||
|
||
except Exception as e:
|
||
error_msg = f"下载失败: {str(e)}"
|
||
print(error_msg)
|
||
if 'temp_file_path' in locals():
|
||
try:
|
||
os.remove(temp_file_path)
|
||
except:
|
||
pass
|
||
raise Exception(error_msg)
|
||
|
||
|
||
def slice_large_image(image_path, slice_size=1024, overlap=100):
|
||
"""切分大图为切片、返回切片数据和位置信息"""
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
raise ValueError(f"无法读取图像: {image_path}")
|
||
|
||
h, w = img.shape[:2]
|
||
step = slice_size - overlap
|
||
num_rows = (h + step - 1) // step
|
||
num_cols = (w + step - 1) // step
|
||
|
||
slices = []
|
||
for i in range(num_rows):
|
||
for j in range(num_cols):
|
||
y1 = i * step
|
||
x1 = j * step
|
||
y2 = min(y1 + slice_size, h)
|
||
x2 = min(x1 + slice_size, w)
|
||
|
||
if y2 - y1 < slice_size:
|
||
y1 = max(0, y2 - slice_size)
|
||
if x2 - x1 < slice_size:
|
||
x1 = max(0, x2 - slice_size)
|
||
|
||
slice_img = img[y1:y2, x1:x2]
|
||
slices.append((x1, y1, slice_img))
|
||
|
||
return slices, (h, w)
|
||
|
||
|
||
def extract_detection_info(result, slice_offset_x, slice_offset_y):
|
||
"""从YOLO OBB结果中提取检测框信息（修正宽高计算）"""
|
||
detections = []
|
||
|
||
if result.obb is not None and len(result.obb) > 0:
|
||
obb_data = result.obb
|
||
obb_xyxy = obb_data.xyxy.cpu().numpy()
|
||
classes = obb_data.cls.cpu().numpy()
|
||
confidences = obb_data.conf.cpu().numpy()
|
||
|
||
for i in range(len(obb_data)):
|
||
x1_slice, y1_slice, x2_slice, y2_slice = obb_xyxy[i]
|
||
# 计算实际宽高（x方向为宽，y方向为高）
|
||
width = x2_slice - x1_slice
|
||
height = y2_slice - y1_slice
|
||
|
||
# 转换为全局坐标
|
||
x1_global = x1_slice + slice_offset_x
|
||
y1_global = y1_slice + slice_offset_y
|
||
|
||
cls_id = int(classes[i])
|
||
confidence = float(confidences[i])
|
||
class_name = result.names[cls_id]
|
||
|
||
detection_info = {
|
||
"type": class_name,
|
||
"size": [int(round(width)), int(round(height))],
|
||
"leftTopPoint": [int(round(x1_global)), int(round(y1_global))],
|
||
"score": round(confidence, 4)
|
||
}
|
||
detections.append(detection_info)
|
||
|
||
return detections
|
||
|
||
|
||
def detect_large_image_from_url(image_url: str, slice_size: int = 1024, overlap: int = 100) -> DetectionResponse:
|
||
"""
|
||
封装后的检测方法：从图片URL处理大图、返回DetectionResponse对象
|
||
"""
|
||
# 动态拼接固定model_path（当前文件同级目录下）
|
||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||
model_path = os.path.join(current_dir, "models", "solor_bracket_model", "best.pt")
|
||
|
||
processing_errors = []
|
||
all_detections = []
|
||
original_size = [0, 0]
|
||
|
||
try:
|
||
# 验证模型文件是否存在
|
||
if not os.path.exists(model_path):
|
||
raise FileNotFoundError(f"模型文件不存在：{model_path}")
|
||
|
||
# 下载图像
|
||
temp_file_path = download_large_file(image_url)
|
||
|
||
try:
|
||
# 切分图像
|
||
slices_info, (h, w) = slice_large_image(temp_file_path, slice_size, overlap)
|
||
original_size = [w, h]
|
||
print(f"完成切片: 共 {len(slices_info)} 个切片")
|
||
|
||
# 加载模型并预测
|
||
model = YOLO(model_path)
|
||
print("开始逐张预测切片...")
|
||
|
||
for i, (x1, y1, slice_img) in enumerate(slices_info, 1):
|
||
print(f"预测第 {i}/{len(slices_info)} 个切片")
|
||
result = model(slice_img, conf=0.5, verbose=False)[0]
|
||
slice_detections = extract_detection_info(result, x1, y1)
|
||
all_detections.extend(slice_detections)
|
||
print(f" 本切片检测到 {len(slice_detections)} 个目标")
|
||
|
||
finally:
|
||
# 确保临时文件删除
|
||
if os.path.exists(temp_file_path):
|
||
try:
|
||
os.remove(temp_file_path)
|
||
print("临时文件已删除")
|
||
except Exception as e:
|
||
error_msg = f"删除临时文件失败: {str(e)}"
|
||
print(error_msg)
|
||
processing_errors.append(error_msg)
|
||
|
||
except Exception as e:
|
||
# 捕获所有异常并记录
|
||
error_msg = str(e)
|
||
processing_errors.append(error_msg)
|
||
print(f"处理异常: {error_msg}")
|
||
|
||
# 构建并返回DetectionResponse对象
|
||
return DetectionResponse(
|
||
hasTarget=1 if len(all_detections) > 0 else 0,
|
||
originalImgSize=original_size,
|
||
targets=all_detections,
|
||
processing_errors=processing_errors
|
||
)
|