feat: add weight to model/

3bda1a9 10 months ago

8.85 kB

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

	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