code_hw_object_v2 / code /file_process_v2.py
zry-research's picture
feat: add weight to model/
3bda1a9
Raw
History Blame
5.32 kB
import os
import cv2
#from paddleocr import PaddleOCRVL
import time
import csv
####
from rule_11 import is_fontsize_percentage_large_image
#from rule_14 import is_text_design_uncoordinated_image
from rule_16 import is_guide_button_with_yolo_image
'''
from rule_2 import is_low_hue_diversity_image
from rule_3 import is_promotion_image
from rule_4 import is_low_value_image
from rule_8 import is_promo_with_yolo_image
from rule_9 import is_subject_too_large_image
from rule_10 import is_outside_safe_area_image
'''
# 假定 ocr 及 is_promotion_image 已初始化/定义
from paddleocr import PaddleOCR
from ultralytics import YOLO
ocr = PaddleOCR(
# ==== ① 文档方向 + 纠偏预处理 ====
use_textline_orientation=False,
use_doc_orientation_classify=False,
use_doc_unwarping=False,
# ==== ② 检测模型 ====
text_detection_model_name="PP-OCRv5_server_det",
text_detection_model_dir="model/paddleocr/det/PP-OCRv5_server_det_infer",
#r"model/paddleocr/det/ch_PP-OCRv3_det_infer_new"
# ==== ③ 识别模型 ====
text_recognition_model_name="PP-OCRv5_server_rec",
text_recognition_model_dir="model/paddleocr/rec/PP-OCRv5_server_rec_infer",
#r"model/paddleocr/rec/ch_PP-OCRv3_rec_infer_new"
# ==== ⑦ 语言 ====
lang="ch"
)
yolo_model = YOLO("model/runs/detect/train8/weights/best.pt")
def main(input_folder):
count_total = 0
count_is_fontsize_percentage_large=0
count_is_text_design_uncoordinated =0
count_is_guide_button_with_yolo=0
'''
count_promo = 0
count_low_value=0
count_low_hue_diversity=0
count_is_promo_with_yolo=0
count_is_subject_too_large=0
count_is_outside_safe_area=0
'''
results_list = []
for filename in os.listdir(input_folder):
if not filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp','.webp','.jfif')):
continue
image_path = os.path.join(input_folder, filename)
img = cv2.imread(image_path)
if img is None:
print(f"[警告] 图片读取失败: {filename}")
continue
count_total += 1
img_result = {
"filename": filename,
"is_fontsize_percentage_large": None,
"is_text_design_uncoordinated": None,
}
'''
'''
is_fontsize_percentage_large= is_fontsize_percentage_large_image(img,ocr)
img_result["is_fontsize_percentage_large"] = is_fontsize_percentage_large
if is_fontsize_percentage_large:
count_is_fontsize_percentage_large += 1
result_text = "文字-占比过大"
else:
result_text = "文字-占比正常"
#print(f"文件: {filename}\t判断: {result_text}")
print(f"文件: {os.path.join(input_folder,filename)}\n判断: {result_text}")
'''
is_text_design_uncoordinated = is_text_design_uncoordinated_image(img,ocr)
img_result["is_text_design_uncoordinated"] = is_text_design_uncoordinated
if is_text_design_uncoordinated:
count_is_text_design_uncoordinated += 1
result_text = "文字搭配不协调"
else:
result_text = "文字搭配协调"
#print(f"文件: {filename}\t判断: {result_text}")
print(f"文件: {os.path.join(input_folder,filename)}\n判断: {result_text}")
is_guide_button_with_yolo = is_guide_button_with_yolo_image(img,model=yolo_model,idx=count_total,ocr=ocr)
img_result["is_text_design_uncoordinated"] = is_guide_button_with_yolo
if is_guide_button_with_yolo:
count_is_guide_button_with_yolo += 1
result_text = "存在引导按钮"
else:
result_text = "不存在"
#print(f"文件: {filename}\t判断: {result_text}")
print(f"文件: {os.path.join(input_folder,filename)}\n判断: {result_text}")
'''
results_list.append(img_result)
# 输出统计信息
if count_total > 0:
percentage = round(count_is_fontsize_percentage_large / count_total * 100, 2)
print(f"\n统计:总图片数={count_total},\
文字占比大={count_is_fontsize_percentage_large}|{round(count_is_fontsize_percentage_large / count_total * 100, 2)}%,\
文字搭配不协调={count_is_text_design_uncoordinated}|{round(count_is_text_design_uncoordinated/ count_total * 100, 2)}%,\
存在引导按钮={count_is_guide_button_with_yolo}|{round(count_is_guide_button_with_yolo/ count_total * 100, 2)}%,\
end")
else:
print("未检测到有效图片。")
output_csv = "image_rule_results.csv"
with open(output_csv, "w", newline='', encoding="utf-8") as csvfile:
fieldnames = [
"filename",
"is_fontsize_percentage_large",
"is_text_design_uncoordinated",
"is_guide_button_with_yolo",
]
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
for row in results_list:
writer.writerow(row)
print(f"\n已保存统计表:{output_csv}")
if __name__ == '__main__':
input_folder = "stage2/hw_v2/文字占比" # 你的图片文件夹路径
main(input_folder)