code_hw_object_v2 / code /file_process_v2.py

feat: add weight to model/

3bda1a9 10 months ago

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)
    

	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)