|
|
#!/usr/bin/env python3
|
|
|
# -*- coding: utf-8 -*-
|
|
|
|
|
|
import os
|
|
|
import sys
|
|
|
import cv2
|
|
|
import argparse
|
|
|
from pathlib import Path
|
|
|
import numpy as np
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
# 添加项目根目录到Python路径
|
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
sys.path.append(script_dir)
|
|
|
|
|
|
# 检查是否可以导入高级检测器
|
|
|
try:
|
|
|
# 导入分析器和高级检测器
|
|
|
from scripts.ship_analyzer import ShipAnalyzer
|
|
|
from utils.advanced_detector import AdvancedShipDetector
|
|
|
ADVANCED_DETECTOR_AVAILABLE = True
|
|
|
except ImportError as e:
|
|
|
print(f"警告：无法导入高级检测器: {e}")
|
|
|
print("将仅使用传统分析器")
|
|
|
from scripts.ship_analyzer import ShipAnalyzer
|
|
|
ADVANCED_DETECTOR_AVAILABLE = False
|
|
|
|
|
|
def analyze_image(image_path, output_dir=None, conf_threshold=0.25, part_conf_threshold=0.3, use_advanced=True):
|
|
|
"""
|
|
|
分析图像中的舰船和部件
|
|
|
|
|
|
Args:
|
|
|
image_path: 图像路径
|
|
|
output_dir: 输出目录
|
|
|
conf_threshold: 检测置信度阈值
|
|
|
part_conf_threshold: 部件置信度阈值
|
|
|
use_advanced: 是否使用高级检测器
|
|
|
"""
|
|
|
print(f"开始分析图像: {image_path}")
|
|
|
|
|
|
# 检查图像是否存在
|
|
|
if not os.path.exists(image_path):
|
|
|
print(f"错误: 图像文件不存在: {image_path}")
|
|
|
return None
|
|
|
|
|
|
# 创建输出目录
|
|
|
if output_dir is not None:
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
|
|
# 根据参数选择使用高级检测器或传统分析器
|
|
|
if use_advanced and ADVANCED_DETECTOR_AVAILABLE:
|
|
|
try:
|
|
|
print("使用高级图像分析器...")
|
|
|
result_img, results = analyze_with_advanced_detector(image_path, output_dir, conf_threshold, part_conf_threshold)
|
|
|
except Exception as e:
|
|
|
print(f"高级分析器出错: {str(e)}")
|
|
|
print("回退到传统分析器...")
|
|
|
# 如果高级分析失败，回退到传统分析器
|
|
|
analyzer = ShipAnalyzer()
|
|
|
results, result_img = analyzer.analyze_image(
|
|
|
image_path,
|
|
|
conf_threshold=conf_threshold,
|
|
|
part_conf_threshold=part_conf_threshold,
|
|
|
save_result=True,
|
|
|
output_dir=output_dir
|
|
|
)
|
|
|
else:
|
|
|
# 使用传统分析器
|
|
|
print("使用传统图像分析器...")
|
|
|
analyzer = ShipAnalyzer()
|
|
|
results, result_img = analyzer.analyze_image(
|
|
|
image_path,
|
|
|
conf_threshold=conf_threshold,
|
|
|
part_conf_threshold=part_conf_threshold,
|
|
|
save_result=True,
|
|
|
output_dir=output_dir
|
|
|
)
|
|
|
|
|
|
# 输出分析结果
|
|
|
if 'ships' in results:
|
|
|
ships = results['ships']
|
|
|
print(f"\n分析完成，检测到 {len(ships)} 个舰船:")
|
|
|
|
|
|
for i, ship in enumerate(ships):
|
|
|
print(f"\n舰船 #{i+1}:")
|
|
|
print(f" 类型: {ship['class_name']}")
|
|
|
print(f" 置信度: {ship['class_confidence']:.2f}")
|
|
|
parts = ship.get('parts', [])
|
|
|
print(f" 检测到 {len(parts)} 个部件:")
|
|
|
|
|
|
# 显示部件信息
|
|
|
for j, part in enumerate(parts):
|
|
|
print(f" 部件 #{j+1}: {part['name']} (置信度: {part['confidence']:.2f})")
|
|
|
else:
|
|
|
# 兼容旧格式
|
|
|
print(f"\n分析完成，检测到 {len(results)} 个舰船:")
|
|
|
for i, ship in enumerate(results):
|
|
|
print(f"\n舰船 #{i+1}:")
|
|
|
print(f" 类型: {ship['class_name']}")
|
|
|
confidence = ship.get('class_confidence', ship.get('confidence', 0.0))
|
|
|
print(f" 置信度: {confidence:.2f}")
|
|
|
parts = ship.get('parts', [])
|
|
|
print(f" 检测到 {len(parts)} 个部件:")
|
|
|
|
|
|
# 显示部件信息
|
|
|
for j, part in enumerate(parts):
|
|
|
part_conf = part.get('confidence', 0.0)
|
|
|
print(f" 部件 #{j+1}: {part['name']} (置信度: {part_conf:.2f})")
|
|
|
|
|
|
# 保存结果图像
|
|
|
if output_dir is not None:
|
|
|
result_path = os.path.join(output_dir, f"analysis_{os.path.basename(image_path)}")
|
|
|
cv2.imwrite(result_path, result_img)
|
|
|
print(f"\n结果图像已保存至: {result_path}")
|
|
|
|
|
|
return result_img
|
|
|
|
|
|
def analyze_with_advanced_detector(image_path, output_dir=None, conf_threshold=0.25, part_conf_threshold=0.3):
|
|
|
"""
|
|
|
使用高级检测器分析图像
|
|
|
|
|
|
Args:
|
|
|
image_path: 图像路径
|
|
|
output_dir: 输出目录
|
|
|
conf_threshold: 检测置信度阈值
|
|
|
part_conf_threshold: 部件置信度阈值
|
|
|
|
|
|
Returns:
|
|
|
result_img: 标注了检测结果的图像
|
|
|
results: 检测结果字典
|
|
|
"""
|
|
|
try:
|
|
|
print("正在加载高级图像分析模型...")
|
|
|
# 初始化高级检测器
|
|
|
detector = AdvancedShipDetector()
|
|
|
except Exception as e:
|
|
|
print(f"高级模型加载失败: {e}")
|
|
|
print("将使用传统计算机视觉方法进行舰船识别")
|
|
|
# 创建一个基本的检测器实例，但不加载模型
|
|
|
detector = AdvancedShipDetector(load_models=False)
|
|
|
|
|
|
# 读取图像
|
|
|
img = cv2.imread(image_path)
|
|
|
if img is None:
|
|
|
raise ValueError(f"无法读取图像: {image_path}")
|
|
|
|
|
|
result_img = img.copy()
|
|
|
h, w = img.shape[:2]
|
|
|
|
|
|
# 使用高级检测器进行对象检测
|
|
|
ships = []
|
|
|
try:
|
|
|
if hasattr(detector, 'detect_ships') and callable(detector.detect_ships):
|
|
|
detected_ships = detector.detect_ships(img, conf_threshold)
|
|
|
if detected_ships and len(detected_ships) > 0:
|
|
|
ships = detected_ships
|
|
|
# 使用检测器返回的图像
|
|
|
if len(detected_ships) > 1 and isinstance(detected_ships[1], np.ndarray):
|
|
|
result_img = detected_ships[1]
|
|
|
ships = detected_ships[0]
|
|
|
else:
|
|
|
print("高级检测器缺少detect_ships方法，使用基本识别")
|
|
|
except Exception as e:
|
|
|
print(f"高级舰船检测失败: {e}")
|
|
|
|
|
|
# 如果没有检测到舰船，使用传统方法尝试识别单个舰船
|
|
|
if not ships:
|
|
|
# 识别舰船类型
|
|
|
ship_type, confidence = detector.identify_ship_type(img)
|
|
|
print(f"高级检测器识别结果: {ship_type}, 置信度: {confidence:.2f}")
|
|
|
|
|
|
# 单个舰船的边界框 - 使用整个图像
|
|
|
padding = int(min(w, h) * 0.05) # 5%的边距
|
|
|
ship_box = (padding, padding, w-padding, h-padding)
|
|
|
|
|
|
# 创建单个舰船对象
|
|
|
ship = {
|
|
|
'id': 1,
|
|
|
'bbox': ship_box,
|
|
|
'class_name': ship_type,
|
|
|
'class_confidence': confidence
|
|
|
}
|
|
|
ships = [ship]
|
|
|
|
|
|
# 在图像上标注舰船信息
|
|
|
cv2.rectangle(result_img, (ship_box[0], ship_box[1]), (ship_box[2], ship_box[3]), (0, 0, 255), 2)
|
|
|
cv2.putText(result_img, f"{ship_type}: {confidence:.2f}",
|
|
|
(ship_box[0]+10, ship_box[1]+30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)
|
|
|
|
|
|
# 为每艘舰船检测部件
|
|
|
processed_ships = []
|
|
|
for i, ship in enumerate(ships):
|
|
|
ship_id = i + 1
|
|
|
ship_box = ship.get('bbox', (0, 0, w, h))
|
|
|
ship_type = ship.get('class_name', '其他舰船')
|
|
|
ship_confidence = ship.get('class_confidence', ship.get('confidence', 0.7))
|
|
|
|
|
|
# 格式化为标准结构
|
|
|
ship_with_parts = {
|
|
|
'id': ship_id,
|
|
|
'bbox': ship_box,
|
|
|
'class_name': ship_type,
|
|
|
'class_confidence': ship_confidence,
|
|
|
'parts': []
|
|
|
}
|
|
|
|
|
|
# 检测舰船部件
|
|
|
try:
|
|
|
parts = detector.detect_ship_parts(img, ship_box, ship_type, part_conf_threshold)
|
|
|
print(f"舰船 #{ship_id} 检测到 {len(parts)} 个部件")
|
|
|
|
|
|
# 为每个部件添加所属舰船ID
|
|
|
for part in parts:
|
|
|
part['ship_id'] = ship_id
|
|
|
ship_with_parts['parts'].append(part)
|
|
|
|
|
|
# 标注部件
|
|
|
part_box = part.get('bbox', (0, 0, 0, 0))
|
|
|
name = part.get('name', '未知部件')
|
|
|
conf = part.get('confidence', 0.0)
|
|
|
|
|
|
# 绘制部件边界框
|
|
|
cv2.rectangle(result_img,
|
|
|
(int(part_box[0]), int(part_box[1])),
|
|
|
(int(part_box[2]), int(part_box[3])),
|
|
|
(0, 255, 0), 2)
|
|
|
|
|
|
# 添加部件标签
|
|
|
label = f"{name}: {conf:.2f}"
|
|
|
cv2.putText(result_img, label,
|
|
|
(int(part_box[0]), int(part_box[1])-5),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
except Exception as e:
|
|
|
print(f"部件检测失败: {e}")
|
|
|
|
|
|
processed_ships.append(ship_with_parts)
|
|
|
|
|
|
# 构建结果数据结构
|
|
|
results = {
|
|
|
'ships': processed_ships
|
|
|
}
|
|
|
|
|
|
# 保存结果图像
|
|
|
if output_dir is not None:
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
result_path = os.path.join(output_dir, f"analysis_{os.path.basename(image_path)}")
|
|
|
cv2.imwrite(result_path, result_img)
|
|
|
print(f"结果图像已保存至: {result_path}")
|
|
|
|
|
|
return result_img, results
|
|
|
|
|
|
def main():
|
|
|
parser = argparse.ArgumentParser(description="舰船图像分析工具")
|
|
|
parser.add_argument("image_path", help="需要分析的舰船图像路径")
|
|
|
parser.add_argument("--output", "-o", help="分析结果输出目录", default="results")
|
|
|
parser.add_argument("--conf", "-c", type=float, default=0.25, help="检测置信度阈值")
|
|
|
parser.add_argument("--part-conf", "-pc", type=float, default=0.3, help="部件检测置信度阈值")
|
|
|
parser.add_argument("--show", action="store_true", help="显示分析结果图像")
|
|
|
parser.add_argument("--traditional", action="store_true", help="使用传统分析器而非高级分析器")
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
try:
|
|
|
# 分析图像
|
|
|
result_img = analyze_image(
|
|
|
args.image_path,
|
|
|
output_dir=args.output,
|
|
|
conf_threshold=args.conf,
|
|
|
part_conf_threshold=args.part_conf,
|
|
|
use_advanced=not args.traditional
|
|
|
)
|
|
|
|
|
|
# 显示结果图像
|
|
|
if args.show and result_img is not None:
|
|
|
cv2.imshow("分析结果", result_img)
|
|
|
cv2.waitKey(0)
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f"分析过程中出错: {str(e)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main() |
