feat: add weight to model/

code/file_process_v2.py

@@ -0,0 +1,153 @@
+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)
+    

code/rule_16.py

@@ -0,0 +1,250 @@
+import os, json
+import numpy as np
+from pathlib import Path
+from typing import Iterable, Optional, Union
+from ultralytics import YOLO
+import re
+
+
+promotion_keywords = [
+    r"领", r"申请", r"立即",r"马上",r"即刻",r"立即",r"下载",
+    r"一定",r"现发",
+    r"大牌美食",
+]
+extra_keywords = [
+    r"立即领取",r"点击领取",r"申请我的额度",r"免费观看",
+    r"火热选购",
+    r"立即前往",r"立即下载",r"立即投保",r"立即参与",r"立即解救",r"立即领取",r"立即抢购",r"立即购买",r"立即签到",
+    r"测一测",r"一定要买",r"现摘现发",
+    r"登录",r"上滑",
+
+]
+compiled_patterns = [re.compile(p, flags=re.IGNORECASE) for p in (promotion_keywords + extra_keywords)]
+compiled_extra_patterns = [re.compile(p, flags=re.IGNORECASE) for p in extra_keywords]
+
+def extract_texts(ocr_result) -> str:
+    # 支持 str / [str] / dict / list[dict/...]
+    texts = []
+    texts = ocr_result[0]['rec_texts']
+    scores = ocr_result[0]['rec_scores']
+    coordss = ocr_result[0]['rec_polys']
+    ocr_result = [
+    (pts, (txt, conf))
+    for pts, txt, conf in zip(coordss, texts, scores)
+    ]
+
+    for res in ocr_result:
+        coords, (text, confidence) = res
+        if confidence > 0.5:
+            texts.append(text)
+
+
+    merged = " ".join([t for t in texts if t]).strip()
+    merged = re.sub(r"\s+", "", merged)
+    print(merged)
+    return merged
+
+def merge_by_rows(ocr_result, y_threshold=40):
+    blocks = []
+    texts = ocr_result[0]['rec_texts']
+    scores = ocr_result[0]['rec_scores']
+    coordss = ocr_result[0]['rec_polys']
+    ocr_result = [
+    (pts, (txt, conf))
+    for pts, txt, conf in zip(coordss, texts, scores)
+    ]
+
+    for pts, (txt, conf) in ocr_result:
+        if conf < 0.5 or not txt.strip():
+            continue
+        ys = [p[1] for p in pts]
+        xs = [p[0] for p in pts]
+        blocks.append({
+            "text": txt.strip(),
+            "x_min": min(xs),
+            "y_max": max(ys)
+        })
+    if not blocks:
+        return ""
+    blocks.sort(key=lambda b: b["y_max"])
+    lines, cur = [], [blocks[0]]
+    for blk in blocks[1:]:
+        if abs(blk["y_max"] - cur[0]["y_max"]) <= y_threshold:
+            cur.append(blk)
+        else:
+            lines.append(cur)
+            cur = [blk]
+    lines.append(cur)
+    merged = []
+    for line in lines:
+        line.sort(key=lambda b: b["x_min"])
+        merged.append("".join(b["text"] for b in line))
+    return " ".join(merged)
+
+# ========== 依赖函数（全部保留，原样引入即可） ==========
+def is_guide_button_with_yolo_image(
+    image_np: np.ndarray,
+    model: Union[str, YOLO],
+    ocr,
+    out_dir: Union[str, Path] = "yolo_outputs",
+    name: str = "predict",
+    imgsz: int = 1024,
+    device: Union[int, str] = 1,
+    iou: float = 0.45,
+    max_det: int = 300,
+    half: bool = False,
+    verbose: bool = True,
+    idx: Optional[Union[int, str]] = None   # 新增：指定保存文件名
+) -> bool:
+    """
+    用 YOLO 对单张 numpy 图像进行检测，并保存可视化与 JSON。
+    返回：是否包含促销元素（True/False）
+    """
+    class_thresholds = {
+        0: 1.0,  # 优惠券
+        1: 0.75,  # 促销贴纸
+        2: 1.0,  # 元宝
+        3: 1.0,  # 大礼物盒子
+        4: 0.95,  # 爆炸
+        5: 0.95,  # 礼物盒子
+        6: 0.95,  # 福袋
+        7: 0.95,  # 红包
+        8: 0.95,  # 金币
+        9: 0.95,  # 金币簇
+        10: 0.35,  # 长条促销贴纸
+        11: 0.35,  # 长条贴纸
+    }
+
+    # 加载模型
+    if isinstance(model, str):
+        model = YOLO(model)
+
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    results = model.predict(
+        source=[image_np],
+        imgsz=imgsz,
+        conf=0.25,
+        iou=iou,
+        device=device,
+        max_det=max_det,
+        half=half,
+        save=True,
+        project=str(out_dir),
+        name=name,
+        exist_ok=True,
+        verbose=verbose
+    )
+
+    promo_found = False
+    det_list = []
+    names_dict = getattr(model, "names", None)
+
+    for result in results:
+        save_dir = Path(result.save_dir)
+        if result.boxes is None:
+            continue
+        boxes = result.boxes.xyxy.cpu().numpy()
+        classes = result.boxes.cls.cpu().numpy().astype(int)
+        scores = result.boxes.conf.cpu().numpy()
+
+        if verbose:
+            print("—— detections ——")
+        
+        h, w = image_np.shape[:2]
+        total_img_area = h * w  # 图像总面积
+
+        for (x1, y1, x2, y2), cid, score in zip(boxes, classes, scores):
+            cname = names_dict.get(int(cid), str(int(cid))) if isinstance(names_dict, dict) else str(int(cid))
+            if verbose:
+                print(f"cls={cname}({cid})  conf={score:.3f}")
+            det_list.append([cname])
+            det_list.append([float(x1), float(y1), float(x2), float(y2), int(cid), float(score)])
+
+            th = class_thresholds.get(cid, None)
+            if th is None:
+                continue
+
+            # ================= 修改部分开始 =================
+            
+            # 处理非长条类型的类别 (即 0-9)
+            if cid not in (10, 11):
+                if cid == 1:
+                    # 计算当前检测框的面积
+                    box_area = (x2 - x1) * (y2 - y1)
+                    # 计算占比
+                    ratio = box_area / total_img_area if total_img_area > 0 else 0
+                    
+                    # 逻辑：如果是类别1，置信度达标 且 面积占比小于 0.01，判定为True
+                    if score >= th and ratio < 0.01:
+                        if verbose:
+                            print(f"   [Match] Class 1 (Area Ratio: {ratio:.4f} < 0.01)")
+                        promo_found = True
+                else:
+                    # 其他类别 (0, 2-9) 保持原逻辑：只看置信度
+                    if score >= th:
+                        promo_found = True
+                continue
+            
+            # ================= 修改部分结束 =================
+
+            # 下面是处理类别 10, 11 的逻辑 (原代码保持不变)
+            if cid in (10,11) and score >= 0.875:
+                promo_found = True
+                continue
+
+            if score >= th:
+                if ocr is not None:
+                    # 裁剪检测框区域
+                    xi1, yi1 = max(0, int(x1)), max(0, int(y1))
+                    xi2, yi2 = min(w - 1, int(x2)), min(h - 1, int(y2))
+                    if xi2 > xi1 and yi2 > yi1:
+                        crop = image_np[yi1:yi2, xi1:xi2]
+                        try:
+                            ocr_result = ocr.ocr(crop)
+                            # text_merged = extract_texts(ocr_result)
+                            text_merged = merge_by_rows(ocr_result)
+                            print(text_merged)
+                            # 正则匹配：命中任一关键词则判定为促销
+                            if any(p.search(text_merged) for p in compiled_patterns):
+                                if verbose:
+                                    print(f"[OCR-MATCH] {cname} 命中促销关键词")
+                                promo_found = True
+                                continue
+                        except Exception as e:
+                            if verbose:
+                                print(f"[OCR-ERROR] {e}")
+                            # OCR 失败时不改变 promo_found，保持谨慎
+
+        # —— 重命名保存的图片 —— 
+        if idx is not None:
+            default_path = save_dir / "image0.jpg"
+            target_path = save_dir / f"{idx}.jpg"
+            if default_path.exists():
+                default_path.rename(target_path)
+                if verbose:
+                    print(f"[保存] 可视化检测图: {target_path}")
+
+        # 保存检测结果 JSON
+        json_path = save_dir / f"detections_{idx if idx is not None else '0'}.json"
+        with open(json_path, "w", encoding="utf-8") as f:
+            json.dump(det_list, f, ensure_ascii=False, indent=2)
+        if verbose:
+            print(f"[保存] 检测数组: {json_path}")
+
+    # ==== 检测框全部没命中，再做一次全局OCR，匹配extra_keywords ====
+    if ocr is not None and not promo_found:
+        try:
+            global_ocr_result = ocr.ocr(image_np)
+            global_text = merge_by_rows(global_ocr_result)
+            print("[全局OCR]", global_text)
+            if any(p.search(global_text) for p in compiled_extra_patterns):
+                if verbose:
+                    print("[全局OCR-MATCH] 命中extra_keywords促销关键词")
+                return True
+        except Exception as e:
+            if verbose:
+                print(f"[全局OCR-ERROR] {e}")
+
+    return promo_found