|
|
@@ -1,6 +1,7 @@
|
|
|
import numpy as np
|
|
|
import json
|
|
|
import torch
|
|
|
+import datetime
|
|
|
import torch.nn as nn
|
|
|
import cv2
|
|
|
from pathlib import Path
|
|
|
@@ -260,51 +261,14 @@ class FryBisenetV2Predictor:
|
|
|
with open(json_path, 'w', encoding='utf-8') as f:
|
|
|
json.dump(json_result, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
- def predict_from_image(self, img_bgr: np.ndarray) -> Dict:
|
|
|
- """
|
|
|
- 直接从解码后的图像数据（numpy数组）进行预测。
|
|
|
-
|
|
|
- Args:
|
|
|
- img_bgr: BGR格式的图像，作为一个numpy数组。
|
|
|
-
|
|
|
- Returns:
|
|
|
- 预测结果字典。
|
|
|
- """
|
|
|
- # 检查通道数是否匹配
|
|
|
- shape = img_bgr.shape
|
|
|
- image_channel = shape[2] if len(shape) == 3 else 1
|
|
|
- if image_channel != self.input_channels:
|
|
|
- raise ValueError(
|
|
|
- f"输入图片的通道数和模型不匹配：image_channel：{image_channel}，input_channels：{self.input_channels}")
|
|
|
-
|
|
|
- # 获取原始图片尺寸
|
|
|
- height, width = img_bgr.shape[:2]
|
|
|
- originImgSize = {'width': width, 'height': height}
|
|
|
-
|
|
|
- # 预处理
|
|
|
- imgTensor_CHW_norm = predict_preprocess(img_bgr, self.imgSize_train_dict)
|
|
|
-
|
|
|
- # 预测
|
|
|
- ansImgDict = self._predict_tensor(imgTensor_CHW_norm)
|
|
|
-
|
|
|
- # 创建结果
|
|
|
- per_img_seg_result = create_result_singleImg(
|
|
|
- self.seg_class_dict,
|
|
|
- ansImgDict,
|
|
|
- originImgSize,
|
|
|
- self.imgSize_train_dict,
|
|
|
- confidence=self.confidence
|
|
|
- )
|
|
|
-
|
|
|
- return per_img_seg_result
|
|
|
-
|
|
|
- def predict_single_image(self,
|
|
|
- img_path: str,
|
|
|
- save_visualization: bool = True,
|
|
|
- save_json: bool = True,
|
|
|
- answer_json_dir: Optional[str] = None,
|
|
|
- input_channels=3
|
|
|
- ) -> Dict:
|
|
|
+ def predict_single_image_np(self,
|
|
|
+ img_bgr: np.ndarray,
|
|
|
+ image_path_str: str = None,
|
|
|
+ save_visualization: bool = True,
|
|
|
+ save_json: bool = True,
|
|
|
+ answer_json_dir_str: Optional[str] = None,
|
|
|
+ input_channels=3
|
|
|
+ ) -> Dict:
|
|
|
"""
|
|
|
预测单张图片
|
|
|
|
|
|
@@ -312,20 +276,20 @@ class FryBisenetV2Predictor:
|
|
|
img_path: 图片路径
|
|
|
save_visualization: 是否保存可视化结果
|
|
|
save_json: 是否保存JSON结果
|
|
|
- answer_json_dir: JSON结果保存目录
|
|
|
+ answer_json_dir_str: JSON结果保存目录
|
|
|
|
|
|
Returns:
|
|
|
预测结果字典
|
|
|
"""
|
|
|
+ if image_path_str is None:
|
|
|
+ timestamp = datetime.now().strftime('%y%m%d_%H%M%S_%f')
|
|
|
+ image_path_real_str = f"{timestamp}.jpg"
|
|
|
+ else:
|
|
|
+ image_path_real_str = str(image_path_str)
|
|
|
|
|
|
- img_path_obj = Path(img_path).resolve()
|
|
|
- img_path_parent_obj = img_path_obj.parent
|
|
|
- answer_json_dir_obj = Path(answer_json_dir).resolve()
|
|
|
+ img_path_real_obj = Path(image_path_real_str).resolve()
|
|
|
|
|
|
- # 读取图片
|
|
|
- img_bgr = cv2.imread(img_path)
|
|
|
- if img_bgr is None:
|
|
|
- raise ValueError(f"无法读取图片：{img_path}")
|
|
|
+ answer_json_dir_str_obj = Path(answer_json_dir_str).resolve()
|
|
|
|
|
|
shape = img_bgr.shape
|
|
|
image_channel = shape[2]
|
|
|
@@ -355,22 +319,27 @@ class FryBisenetV2Predictor:
|
|
|
# 预测
|
|
|
ansImgDict = self._predict_tensor(imgTensor_CHW_norm)
|
|
|
|
|
|
+ image_path_name = str(img_path_real_obj.name)
|
|
|
+
|
|
|
# 创建结果
|
|
|
per_img_seg_result = create_result_singleImg(
|
|
|
self.seg_class_dict,
|
|
|
ansImgDict,
|
|
|
originImgSize,
|
|
|
self.imgSize_train_dict,
|
|
|
- confidence=self.confidence
|
|
|
+ confidence=self.confidence,
|
|
|
)
|
|
|
|
|
|
# 保存JSON结果
|
|
|
- if save_json and answer_json_dir:
|
|
|
- json_dir = Path(answer_json_dir)
|
|
|
+ if save_json and answer_json_dir_str:
|
|
|
+ json_dir = Path(answer_json_dir_str)
|
|
|
json_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
+ if image_path_str is None:
|
|
|
+ cv2.imwrite(image_path_real_str, img_bgr)
|
|
|
+
|
|
|
# 获取图片文件名（不含扩展名）
|
|
|
- img_name = Path(img_path).stem
|
|
|
+ img_name = Path(img_path_real_obj).stem
|
|
|
json_path = json_dir / f"{img_name}.json"
|
|
|
|
|
|
self._save_result_json(per_img_seg_result, json_path)
|
|
|
@@ -379,18 +348,123 @@ class FryBisenetV2Predictor:
|
|
|
# 保存可视化结果
|
|
|
if save_visualization:
|
|
|
result_img = process_detection_result(img_bgr, per_img_seg_result, self.label_colors_dict)
|
|
|
- output_path = str(answer_json_dir_obj / f"{Path(img_path).stem}_result.jpg")
|
|
|
+ output_path = str(answer_json_dir_str_obj / f"{img_path_real_obj.name}")
|
|
|
|
|
|
cv2.imwrite(output_path, result_img)
|
|
|
fry_algo_print("成功", f"可视化结果已保存到：{output_path}")
|
|
|
|
|
|
return per_img_seg_result
|
|
|
|
|
|
+ def _save_result_json(self, result: Dict, json_path: Path):
|
|
|
+ """
|
|
|
+ 保存预测结果为JSON文件
|
|
|
+
|
|
|
+ Args:
|
|
|
+ result: 预测结果字典
|
|
|
+ json_path: JSON文件保存路径
|
|
|
+ """
|
|
|
+ # 将numpy数组转换为可序列化的格式
|
|
|
+ json_result = {}
|
|
|
+
|
|
|
+ for key, value in result.items():
|
|
|
+ if isinstance(value, np.ndarray):
|
|
|
+ json_result[key] = value.tolist()
|
|
|
+ elif isinstance(value, dict):
|
|
|
+ json_result[key] = {}
|
|
|
+ for k, v in value.items():
|
|
|
+ if isinstance(v, np.ndarray):
|
|
|
+ json_result[key][k] = v.tolist()
|
|
|
+ else:
|
|
|
+ json_result[key][k] = v
|
|
|
+ else:
|
|
|
+ json_result[key] = value
|
|
|
+
|
|
|
+ with open(json_path, 'w', encoding='utf-8') as f:
|
|
|
+ json.dump(json_result, f, ensure_ascii=False, indent=2)
|
|
|
+
|
|
|
+ def predict_from_image(self, img_bgr: np.ndarray) -> Dict:
|
|
|
+ """
|
|
|
+ 直接从解码后的图像数据（numpy数组）进行预测。
|
|
|
+
|
|
|
+ Args:
|
|
|
+ img_bgr: BGR格式的图像，作为一个numpy数组。
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 预测结果字典。
|
|
|
+ """
|
|
|
+ # 检查通道数是否匹配
|
|
|
+ shape = img_bgr.shape
|
|
|
+ image_channel = shape[2] if len(shape) == 3 else 1
|
|
|
+ if image_channel != self.input_channels:
|
|
|
+ raise ValueError(
|
|
|
+ f"输入图片的通道数和模型不匹配：image_channel：{image_channel}，input_channels：{self.input_channels}")
|
|
|
+
|
|
|
+ # 获取原始图片尺寸
|
|
|
+ height, width = img_bgr.shape[:2]
|
|
|
+ originImgSize = {'width': width, 'height': height}
|
|
|
+
|
|
|
+ # 预处理
|
|
|
+ imgTensor_CHW_norm = predict_preprocess(img_bgr, self.imgSize_train_dict)
|
|
|
+
|
|
|
+ # 预测
|
|
|
+ ansImgDict = self._predict_tensor(imgTensor_CHW_norm)
|
|
|
+
|
|
|
+ # 创建结果
|
|
|
+ per_img_seg_result = create_result_singleImg(
|
|
|
+ self.seg_class_dict,
|
|
|
+ ansImgDict,
|
|
|
+ originImgSize,
|
|
|
+ self.imgSize_train_dict,
|
|
|
+ confidence=self.confidence
|
|
|
+ )
|
|
|
+
|
|
|
+ return per_img_seg_result
|
|
|
+
|
|
|
+ def predict_single_image(self,
|
|
|
+ img_path: str,
|
|
|
+ save_visualization: bool = True,
|
|
|
+ save_json: bool = True,
|
|
|
+ answer_json_dir_str: Optional[str] = None,
|
|
|
+ input_channels=3
|
|
|
+ ) -> Dict:
|
|
|
+ """
|
|
|
+ 预测单张图片
|
|
|
+
|
|
|
+ Args:
|
|
|
+ img_path: 图片路径
|
|
|
+ save_visualization: 是否保存可视化结果
|
|
|
+ save_json: 是否保存JSON结果
|
|
|
+ answer_json_dir_str: JSON结果保存目录
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 预测结果字典
|
|
|
+ """
|
|
|
+
|
|
|
+ img_path_obj = Path(img_path).resolve()
|
|
|
+ img_path_parent_obj = img_path_obj.parent
|
|
|
+ answer_json_dir_str_obj = Path(answer_json_dir_str).resolve()
|
|
|
+
|
|
|
+ # 读取图片
|
|
|
+ img_bgr = cv2.imread(str(img_path_obj))
|
|
|
+ if img_bgr is None:
|
|
|
+ raise ValueError(f"无法读取图片：{img_path}")
|
|
|
+
|
|
|
+ per_img_seg_result = self.predict_single_image_np(
|
|
|
+ img_bgr=img_bgr,
|
|
|
+ image_path_str=str(img_path_obj),
|
|
|
+ save_visualization=save_visualization,
|
|
|
+ save_json=save_json,
|
|
|
+ answer_json_dir_str=answer_json_dir_str,
|
|
|
+ input_channels=input_channels
|
|
|
+ )
|
|
|
+
|
|
|
+ return per_img_seg_result
|
|
|
+
|
|
|
def predict_batch(self,
|
|
|
img_paths: List[str],
|
|
|
save_visualization: bool = True,
|
|
|
save_json: bool = True,
|
|
|
- answer_json_dir: Optional[str] = None,
|
|
|
+ answer_json_dir_str: Optional[str] = None,
|
|
|
input_channels=3
|
|
|
) -> List[Dict]:
|
|
|
"""
|
|
|
@@ -400,18 +474,18 @@ class FryBisenetV2Predictor:
|
|
|
img_paths: 图片路径列表
|
|
|
save_visualization: 是否保存可视化结果
|
|
|
save_json: 是否保存JSON结果
|
|
|
- answer_json_dir: JSON结果保存目录
|
|
|
+ answer_json_dir_str: JSON结果保存目录
|
|
|
output_dir: 可视化结果保存目录
|
|
|
|
|
|
Returns:
|
|
|
所有图片的预测结果列表
|
|
|
"""
|
|
|
|
|
|
- answer_json_dir_obj = Path(answer_json_dir).resolve()
|
|
|
+ answer_json_dir_str_obj = Path(answer_json_dir_str).resolve()
|
|
|
|
|
|
results = []
|
|
|
|
|
|
- Path(answer_json_dir).mkdir(parents=True, exist_ok=True)
|
|
|
+ Path(answer_json_dir_str).mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
# 批量处理
|
|
|
for i, img_path in enumerate(img_paths):
|
|
|
@@ -444,18 +518,22 @@ class FryBisenetV2Predictor:
|
|
|
# 预测
|
|
|
ansImgDict = self._predict_tensor(imgTensor_CHW_norm)
|
|
|
|
|
|
+ img_path_obj = Path(img_path).resolve()
|
|
|
+ image_path_name = str(img_path_obj.name)
|
|
|
+
|
|
|
# 创建结果
|
|
|
per_img_seg_result = create_result_singleImg(
|
|
|
self.seg_class_dict,
|
|
|
ansImgDict,
|
|
|
originImgSize,
|
|
|
self.imgSize_train_dict,
|
|
|
- confidence=self.confidence
|
|
|
+ confidence=self.confidence,
|
|
|
+ image_path_name=image_path_name
|
|
|
)
|
|
|
|
|
|
# 保存JSON结果
|
|
|
- if save_json and answer_json_dir:
|
|
|
- json_dir = Path(answer_json_dir)
|
|
|
+ if save_json and answer_json_dir_str:
|
|
|
+ json_dir = Path(answer_json_dir_str)
|
|
|
json_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
img_name = Path(img_path).stem
|
|
|
@@ -465,8 +543,7 @@ class FryBisenetV2Predictor:
|
|
|
# 保存可视化结果
|
|
|
if save_visualization:
|
|
|
result_img = process_detection_result(img_bgr, per_img_seg_result, self.label_colors_dict)
|
|
|
-
|
|
|
- output_path = answer_json_dir_obj / f"{Path(img_path).stem}_result.jpg"
|
|
|
+ output_path = answer_json_dir_str_obj / f"{Path(img_path).name}"
|
|
|
|
|
|
cv2.imwrite(str(output_path), result_img)
|
|
|
|
|
|
@@ -483,14 +560,14 @@ class FryBisenetV2Predictor:
|
|
|
def main():
|
|
|
"""使用示例"""
|
|
|
# 配置参数
|
|
|
- pth_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\Model\outer_box.pth"
|
|
|
+ pth_path = r"segmentation_bisenetv2.pth"
|
|
|
input_channels = 3
|
|
|
|
|
|
real_seg_class_dict = {1: 'outer_box'}
|
|
|
|
|
|
# 为不同类别设置不同颜色（可选）
|
|
|
label_colors_dict = {
|
|
|
- 'outer_box': (225, 0, 0),
|
|
|
+ 'outer_box': (255, 0, 0),
|
|
|
}
|
|
|
|
|
|
imgSize_train_dict = {'width': 1280, 'height': 1280}
|
|
|
@@ -508,14 +585,14 @@ def main():
|
|
|
|
|
|
# 单张图片预测
|
|
|
print("=== 单张图片预测 ===")
|
|
|
- now_img_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img.png"
|
|
|
- answer_json_dir = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\outer"
|
|
|
+ now_img_path = r"input_output\images\coaxis_0008.jpg"
|
|
|
+ answer_json_dir_str = r"input_output\images_answer_json_dir_str"
|
|
|
|
|
|
result = predictor.predict_single_image(
|
|
|
img_path=now_img_path,
|
|
|
save_visualization=True,
|
|
|
save_json=True,
|
|
|
- answer_json_dir=answer_json_dir,
|
|
|
+ answer_json_dir_str=answer_json_dir_str,
|
|
|
input_channels=input_channels,
|
|
|
)
|
|
|
|
|
|
@@ -532,14 +609,14 @@ def main():
|
|
|
# img_paths=img_paths,
|
|
|
# save_visualization=True,
|
|
|
# save_json=True,
|
|
|
- # answer_json_dir=answer_json_dir,
|
|
|
+ # answer_json_dir_str=answer_json_dir_str,
|
|
|
# input_channels=input_channels,
|
|
|
# )
|
|
|
|
|
|
|
|
|
def _test_pokemon_inner_box():
|
|
|
# 配置参数
|
|
|
- pth_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\Model\inner_box.pth"
|
|
|
+ pth_path = r"E:\_250807_训练好的导出的模型\_250808_1043_宝可梦内框训练效果还行\pth_and_images\segmentation_bisenetv2.pth"
|
|
|
input_channels = 3
|
|
|
|
|
|
real_seg_class_dict = {1: 'inner_box'}
|
|
|
@@ -564,20 +641,20 @@ def _test_pokemon_inner_box():
|
|
|
|
|
|
# 单张图片预测
|
|
|
print("=== 单张图片预测 ===")
|
|
|
- now_img_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img.png"
|
|
|
- answer_json_dir = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\inner"
|
|
|
+ now_img_path = r"E:\_250807_训练好的导出的模型\_250808_1043_宝可梦内框训练效果还行\pth_and_images\images\diff_big_00065.jpg"
|
|
|
+ answer_json_dir_str = r"E:\_250807_训练好的导出的模型\_250808_1043_宝可梦内框训练效果还行\pth_and_images\images_answer"
|
|
|
|
|
|
result = predictor.predict_single_image(
|
|
|
img_path=now_img_path,
|
|
|
save_visualization=True,
|
|
|
save_json=True,
|
|
|
- answer_json_dir=answer_json_dir
|
|
|
+ answer_json_dir_str=answer_json_dir_str
|
|
|
)
|
|
|
|
|
|
|
|
|
def _test_pokemon_back_edge():
|
|
|
# 配置参数
|
|
|
- pth_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\Model\back_defect.pth"
|
|
|
+ pth_path = r"E:\_250807_训练好的导出的模型\_250811_1104_宝可梦背面边角\pth_and_images\segmentation_bisenetv2.pth"
|
|
|
input_channels = 3
|
|
|
|
|
|
real_seg_class_dict = {
|
|
|
@@ -607,60 +684,19 @@ def _test_pokemon_back_edge():
|
|
|
|
|
|
# 单张图片预测
|
|
|
print("=== 单张图片预测 ===")
|
|
|
- now_img_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img_2.png"
|
|
|
- answer_json_dir = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\defect"
|
|
|
+ now_img_path = r"E:\_250807_训练好的导出的模型\_250811_1104_宝可梦背面边角\pth_and_images\images\split\Pokémon_back_for_Edge_0001_bottom_grid_r0_c0.jpg"
|
|
|
+ answer_json_dir_str = r"E:\_250807_训练好的导出的模型\_250811_1104_宝可梦背面边角\pth_and_images\images_answer"
|
|
|
|
|
|
result = predictor.predict_single_image(
|
|
|
img_path=now_img_path,
|
|
|
save_visualization=True,
|
|
|
save_json=True,
|
|
|
- answer_json_dir=answer_json_dir,
|
|
|
+ answer_json_dir_str=answer_json_dir_str,
|
|
|
input_channels=input_channels
|
|
|
)
|
|
|
|
|
|
-def _test_pokemon_no_reflect_front():
|
|
|
- # 配置参数
|
|
|
- pth_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\Model\no_reflect_front_defect.pth"
|
|
|
- input_channels = 3
|
|
|
-
|
|
|
- real_seg_class_dict = {
|
|
|
- 1: 'scratch',
|
|
|
- 2: 'pit',
|
|
|
- 3: 'stain'
|
|
|
- }
|
|
|
-
|
|
|
- # 为不同类别设置不同颜色（可选）
|
|
|
- # label_colors_dict = {
|
|
|
- # 'outer_box': (255, 0, 0),
|
|
|
- # }
|
|
|
-
|
|
|
- imgSize_train_dict = {'width': 512, 'height': 512}
|
|
|
- confidence = 0.5
|
|
|
-
|
|
|
- # 创建预测器
|
|
|
- predictor = FryBisenetV2Predictor(
|
|
|
- pth_path=pth_path,
|
|
|
- real_seg_class_dict=real_seg_class_dict,
|
|
|
- imgSize_train_dict=imgSize_train_dict,
|
|
|
- confidence=confidence,
|
|
|
- input_channels=input_channels,
|
|
|
- )
|
|
|
-
|
|
|
- # 单张图片预测
|
|
|
- print("=== 单张图片预测 ===")
|
|
|
- now_img_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img_1.png"
|
|
|
- answer_json_dir = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\no_reflect_front_defect"
|
|
|
-
|
|
|
- result = predictor.predict_single_image(
|
|
|
- img_path=now_img_path,
|
|
|
- save_visualization=True,
|
|
|
- save_json=True,
|
|
|
- answer_json_dir=answer_json_dir,
|
|
|
- input_channels=input_channels
|
|
|
- )
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
- # main()
|
|
|
- # _test_pokemon_inner_box()
|
|
|
- # _test_pokemon_back_edge()
|
|
|
- _test_pokemon_no_reflect_front()
|
|
|
+ main()
|
|
|
+ # test_pokemon_inner_box()
|
|
|
+ # test_pokemon_back_edge()
|
AnlaAnla