97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
try:
|
||
from .abdomen_yolo import AbdominalAcupointsDetector as YoloDetector
|
||
from .abdomen_manual import AbdominalAcupointsDetector as ManualDetector
|
||
from .config import Config
|
||
except ImportError:
|
||
from abdomen_yolo import AbdominalAcupointsDetector as YoloDetector
|
||
from abdomen_manual import AbdominalAcupointsDetector as ManualDetector
|
||
from config import Config
|
||
import os
|
||
|
||
class AbdomenDetector:
|
||
"""腹部穴位检测器，整合了YOLO自动检测和手动标注两种方法"""
|
||
|
||
def __init__(self):
|
||
"""
|
||
初始化检测器，同时加载YOLO模型和手动检测器
|
||
"""
|
||
self.yolo = YoloDetector()
|
||
self.manual = ManualDetector()
|
||
|
||
def detect_by_yolo(self, image_path, output_path=None):
|
||
"""
|
||
使用YOLO模型进行自动检测
|
||
|
||
参数:
|
||
image_path (str): 输入图片路径
|
||
output_path (str, optional): 输出图片路径，如果不指定则自动生成
|
||
|
||
返回:
|
||
dict: 包含穴位坐标的字典 {"穴位名称": (x, y), ...}
|
||
"""
|
||
if output_path is None:
|
||
# 从输入图片名生成输出图片名
|
||
base_name = os.path.basename(image_path)
|
||
name, ext = os.path.splitext(base_name)
|
||
output_path = Config.get_output_path(f"{name}_yolo{ext}")
|
||
|
||
return self.yolo.process_image(image_path, output_path)
|
||
|
||
def detect_by_manual(self, image_path, json_data, output_path=None):
|
||
"""
|
||
使用手动标注数据进行检测
|
||
|
||
参数:
|
||
image_path (str): 输入图片路径
|
||
json_data (dict): 包含腹部顶点、脐点和尺寸信息的JSON数据
|
||
output_path (str, optional): 输出图片路径，如果不指定则自动生成
|
||
|
||
返回:
|
||
dict: 包含穴位坐标的字典 {"穴位名称": (x, y), ...}
|
||
"""
|
||
if output_path is None:
|
||
# 从输入图片名生成输出图片名
|
||
base_name = os.path.basename(image_path)
|
||
name, ext = os.path.splitext(base_name)
|
||
output_path = Config.get_output_path(f"{name}_manual{ext}")
|
||
|
||
self.manual.load_data(json_data)
|
||
return self.manual.process_image(image_path, output_path)
|
||
|
||
if __name__ == "__main__":
|
||
# 示例用法
|
||
sample_json = {
|
||
"shapes": [
|
||
{
|
||
"label": "fubu",
|
||
"points": [
|
||
[278, 241],
|
||
[400, 255],
|
||
[396, 301],
|
||
[262, 294]
|
||
],
|
||
"shape_type": "polygon"
|
||
},
|
||
{
|
||
"label": "duqiyan",
|
||
"points": [[341, 276]],
|
||
"shape_type": "point"
|
||
}
|
||
]
|
||
}
|
||
|
||
# 创建检测器实例
|
||
detector = AbdomenDetector()
|
||
|
||
# 测试图片路径
|
||
test_image = Config.get_image_path("color.png")
|
||
|
||
# YOLO检测示例
|
||
yolo_results = detector.detect_by_yolo(test_image)
|
||
print("YOLO Result:")
|
||
print(yolo_results)
|
||
|
||
# 手动标注检测示例
|
||
manual_results = detector.detect_by_manual(test_image, sample_json)
|
||
print("\nManully checked results")
|
||
print(manual_results) |