ensemble_results.py

import os
import re
import json
import glob
import copy
from collections import Counter, defaultdict
from statistics import median

CRITERIA = [
    "Color Harmony",
    "Visual Style Consistency",
    "Sharpness",
    "Light and Shadow Modeling",
    "Creativity and Originality",
    "Exposure Control",
    "Application of Classical Composition Principles",
    "Depth of Field and Layering",
    "Visual Center Stability",
    "Visual Flow Guidance",
    "Structural Support Stability",
    "Appropriateness of Negative Space",
    "Subject Integrity",
]

# LEVEL_ORDER = {"Poor": 0, "Medium": 1, "Good": 2}
# LEVEL_INV = {0: "Poor", 1: "Medium", 2: "Good"}


LEVEL_ORDER = {"Poor": 0, "Medium": 1, "Good": 2}
LEVEL_INV = {0: "A", 1: "B", 2: "C"}


def normalize_level(x):
    if not isinstance(x, str):
        return None
    x = x.strip().lower()
    mp = {
        "poor": "Poor",
        "medium": "Medium",
        "good": "Good",
    }
    return mp.get(x)


def basename_from_item(item):
    img_path = item.get("images", [{}])[0].get("path", "")
    return os.path.basename(img_path)


def parse_response_raw(resp):
    """
    支持：
    - "{\"total_score\": 84}"
    - "41"
    - "{\"criteria\": {...}}"
    - "[\"Medium\", ...]"
    - "{\"answer\": \"C\"}"
    """
    if isinstance(resp, (dict, list, int, float)):
        return resp

    if not isinstance(resp, str):
        return None

    s = resp.strip()

    # 纯数字分数
    if re.fullmatch(r"-?\d+(\.\d+)?", s):
        return float(s)

    # 标准 JSON
    try:
        return json.loads(s)
    except Exception:
        pass

    # 兜底：提取 {...}
    m = re.search(r"\{.*\}", s, flags=re.S)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            pass

    # 兜底：提取 [...]
    m = re.search(r"\[.*\]", s, flags=re.S)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            pass

    return None


def iter_json_or_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        text = f.read().strip()

    if not text:
        return []

    try:
        obj = json.loads(text)
        if isinstance(obj, list):
            return obj
        if isinstance(obj, dict):
            return [obj]
    except Exception:
        pass

    rows = []
    for line in text.splitlines():
        line = line.strip()
        if line:
            rows.append(json.loads(line))
    return rows


def read_all_files(folder, recursive=True):
    pattern = "**/*.json*" if recursive else "*.json*"
    files = sorted(glob.glob(os.path.join(folder, pattern), recursive=recursive))

    rows = []
    for fp in files:
        try:
            rows.extend(iter_json_or_jsonl(fp))
        except Exception as e:
            print(f"[WARN] failed to read {fp}: {e}")
    return rows


def parse_score_folder(folder):
    pred = defaultdict(list)

    for item in read_all_files(folder):
        name = basename_from_item(item)
        data = parse_response_raw(item.get("response", ""))

        score = None

        if isinstance(data, dict):
            score = data.get("total_score")
        elif isinstance(data, (int, float)):
            score = data
        elif isinstance(data, str):
            if re.fullmatch(r"-?\d+(\.\d+)?", data.strip()):
                score = float(data.strip())

        if name and score is not None:
            try:
                score = float(score)
                score = max(0, min(100, score))
                pred[name].append(score)
            except Exception:
                pass

    return pred


def parse_level_folder(folder):
    pred = defaultdict(lambda: defaultdict(list))

    for item in read_all_files(folder):
        name = basename_from_item(item)
        data = parse_response_raw(item.get("response", ""))

        if not name:
            continue

        # 格式1：{"criteria": {"Color Harmony": "Good", ...}}
        if isinstance(data, dict) and isinstance(data.get("criteria"), dict):
            criteria = data["criteria"]
            for c in CRITERIA:
                lv = normalize_level(criteria.get(c))
                if lv:
                    pred[name][c].append(lv)

        # 格式2：["Medium", "Medium", ..., 共13个]
        elif isinstance(data, list):
            for c, lv_raw in zip(CRITERIA, data):
                lv = normalize_level(lv_raw)
                if lv:
                    pred[name][c].append(lv)

    return pred


def parse_reason_folder(folder):
    pred = defaultdict(list)

    for item in read_all_files(folder):
        name = basename_from_item(item)
        data = parse_response_raw(item.get("response", ""))

        ans = None

        if isinstance(data, dict):
            ans = data.get("answer")
        elif isinstance(data, str):
            ans = data

        if name and isinstance(ans, str):
            ans = ans.strip().upper()
            if ans in {"A", "B", "C", "D"}:
                pred[name].append(ans)

    return pred


def majority_vote(values, default=None):
    values = [v for v in values if v is not None]
    if not values:
        return default

    cnt = Counter(values)

    # 平票时按第一次出现顺序
    return max(cnt.keys(), key=lambda x: (cnt[x], -values.index(x)))


def ensemble_scores(score_dicts, method="mean"):
    merged = defaultdict(list)

    # print("merged is", merged)

    for d in score_dicts:
        for name, scores in d.items():
            merged[name].extend(scores)

    out = {}
    for name, scores in merged.items():

        
        if method == "median":
            val = median(scores)
        else:
            val = sum(scores) / len(scores)

        # print("scores are", name, scores, val)

        out[name] = int(round(max(0, min(100, val))))
        # out[name] = int(round(val))

    return out



#####################

LEVEL_SCORE = {
    "Poor": 2.5,
    "Medium": 6.0,
    "Good": 8.5,
}


def score_to_level(score):
    if 0 <= score < 5:
        return "A"
    elif 5 <= score < 7:
        return "B"
    elif 7 <= score <= 10:
        return "C"
    else:
        # 兜底，防止异常值
        score = max(0, min(10, score))
        if score < 5:
            return "A"
        elif score < 7:
            return "B"
        return "C"


def ensemble_levels(level_dicts, method="score_mean"):
    merged = defaultdict(lambda: defaultdict(list))

    for d in level_dicts:
        for name, cd in d.items():
            for c, levels in cd.items():
                merged[name][c].extend(levels)

    out = defaultdict(dict)

    for name, cd in merged.items():
        for c in CRITERIA:
            vals = cd.get(c, [])
            if not vals:
                continue

            if method == "score_mean":
                nums = [LEVEL_SCORE[v] for v in vals if v in LEVEL_SCORE]
                if nums:
                    avg_score = sum(nums) / len(nums)
                    out[name][c] = score_to_level(avg_score)

            elif method == "vote":
                out[name][c] = majority_vote(vals, default="Medium")

            elif method == "ordinal_mean":
                nums = [LEVEL_ORDER[v] for v in vals if v in LEVEL_ORDER]
                if nums:
                    out[name][c] = LEVEL_INV[int(round(sum(nums) / len(nums)))]

    return out



######################


def ensemble_answers(reason_dicts):
    merged = defaultdict(list)

    for d in reason_dicts:
        for name, answers in d.items():
            merged[name].extend(answers)

    return {
        name: majority_vote(answers, default="A")
        for name, answers in merged.items()
    }


def build_submission(
    template_path,
    score_model_folders,
    level_model_folders,
    reason_model_folders,
    output_path,
    score_method="mean",
    level_method="vote",
):
    score_dicts = [parse_score_folder(p) for p in score_model_folders]
    level_dicts = [parse_level_folder(p) for p in level_model_folders]
    reason_dicts = [parse_reason_folder(p) for p in reason_model_folders]

    score_ens = ensemble_scores(score_dicts, method=score_method)
    level_ens = ensemble_levels(level_dicts, method=level_method)
    answer_ens = ensemble_answers(reason_dicts)

    with open(template_path, "r", encoding="utf-8") as f:
        result = json.load(f)

    missing_score = 0
    missing_level = 0
    missing_answer = 0

    for item in result:
        name = item["image_path"]

        if name in score_ens:
            item["total_score"] = score_ens[name]
        else:
            missing_score += 1

        for c in CRITERIA:
            if name in level_ens and c in level_ens[name]:
                item["criteria"][c]["level"] = level_ens[name][c]
            else:
                missing_level += 1

        if name in answer_ens:
            item["answer"] = answer_ens[name]
        else:
            missing_answer += 1

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)

    print(f"Saved to: {output_path}")
    print(f"Images: {len(result)}")
    print(f"Missing score images: {missing_score}")
    print(f"Missing level fields: {missing_level}")
    print(f"Missing answer images: {missing_answer}")


if __name__ == "__main__":

    # from pathlib import Path
    # ROOT = Path("/mnt/shared-storage-user/zhuxiaorong/liyunhao_data/my_code/cvpr26_challenge")


    TEMPLATE_PATH = "track_1_test_demo.json"
    OUTPUT_PATH = "./track_1_test.json"


    SCORE_MODEL_FOLDERS = [
        "./result-score"
    ]
    LEVEL_MODEL_FOLDERS = [
        "./result-level"
    ]
    REASON_MODEL_FOLDERS = [
        "./result-reason"
    ]


    build_submission(
        template_path=TEMPLATE_PATH,
        score_model_folders=SCORE_MODEL_FOLDERS,
        level_model_folders=LEVEL_MODEL_FOLDERS,
        reason_model_folders=REASON_MODEL_FOLDERS,
        output_path=OUTPUT_PATH,
        score_method="mean",
        level_method="score_mean",
    )