|
|
|
|
|
|
| import argparse
|
| import json
|
| import math
|
| import os
|
| import random
|
| from collections import Counter, defaultdict
|
| from dataclasses import dataclass
|
| from typing import Any, Dict, List, Tuple
|
|
|
|
|
| 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",
|
| ]
|
|
|
|
|
| PROMPT_SCORE = (
|
| "You are an aesthetics expert. Evaluate the image on the following 13 criteria:\n"
|
| "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.\n\n"
|
| "For each criterion, output a numeric score in [0.0, 10.0].\n"
|
| "Then output total_score as an integer in [0, 100].\n"
|
| "Return STRICT JSON only.\n"
|
| "JSON format:\n"
|
| "{\n"
|
| " \"criteria\": {\n"
|
| " \"Color Harmony\": {\"score\": 0.0}\n"
|
| " },\n"
|
| " \"total_score\": 0\n"
|
| "}\n\n"
|
| "<image>"
|
| )
|
|
|
|
|
| PROMPT_MULTITASK = (
|
| "You are an aesthetics expert. Evaluate the image on the following 13 criteria:\n"
|
| "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.\n\n"
|
| "For each criterion, output:\n"
|
| "1) score in [0.0, 10.0]\n"
|
| "2) level in {A,B,C} where A:<5, B:[5,7), C:>=7\n"
|
| "Then output total_score as an integer in [0, 100].\n"
|
| "Return STRICT JSON only.\n"
|
| "JSON format:\n"
|
| "{\n"
|
| " \"criteria\": {\n"
|
| " \"Color Harmony\": {\"score\": 0.0, \"level\": \"A|B|C\"}\n"
|
| " },\n"
|
| " \"total_score\": 0\n"
|
| "}\n\n"
|
| "<image>"
|
| )
|
|
|
|
|
| @dataclass
|
| class Sample:
|
| image_path: str
|
| scores: Dict[str, float]
|
| levels: Dict[str, str]
|
| total_score: int
|
| boundary_count: int
|
|
|
|
|
| def clamp_score(x: float) -> float:
|
| return max(0.0, min(10.0, x))
|
|
|
|
|
| def clamp_int(x: int, lo: int, hi: int) -> int:
|
| return max(lo, min(hi, x))
|
|
|
|
|
| def score_to_level(score: float) -> str:
|
| if score < 5.0:
|
| return "A"
|
| if score < 7.0:
|
| return "B"
|
| return "C"
|
|
|
|
|
| def safe_float(x: Any, default: float = 0.0) -> float:
|
| try:
|
| return float(x)
|
| except Exception:
|
| return default
|
|
|
|
|
| def safe_int(x: Any, default: int = 0) -> int:
|
| try:
|
| return int(round(float(x)))
|
| except Exception:
|
| return default
|
|
|
|
|
| def parse_assistant_payload(payload: str) -> Dict[str, Any]:
|
| obj = json.loads(payload)
|
| if not isinstance(obj, dict):
|
| raise ValueError("assistant payload is not dict")
|
| return obj
|
|
|
|
|
| def normalize_record(rec: Dict[str, Any], boundary_margin: float) -> Sample:
|
| assistant = parse_assistant_payload(rec["messages"][1]["content"])
|
| criteria = assistant.get("criteria", {})
|
|
|
| scores: Dict[str, float] = {}
|
| levels: Dict[str, str] = {}
|
| boundary_count = 0
|
|
|
| for c in CRITERIA:
|
| v = (criteria.get(c, {}) or {}).get("score", 0.0)
|
| s = round(clamp_score(safe_float(v, 0.0)), 1)
|
| lv = score_to_level(s)
|
| scores[c] = s
|
| levels[c] = lv
|
| if abs(s - 5.0) <= boundary_margin or abs(s - 7.0) <= boundary_margin:
|
| boundary_count += 1
|
|
|
| t = assistant.get("total_score", 0)
|
| total_score = clamp_int(safe_int(t, 0), 0, 100)
|
|
|
| img = rec.get("images", [""])
|
| image_path = str(img[0]) if isinstance(img, list) and img else ""
|
|
|
| return Sample(
|
| image_path=image_path,
|
| scores=scores,
|
| levels=levels,
|
| total_score=total_score,
|
| boundary_count=boundary_count,
|
| )
|
|
|
|
|
| def build_score_record(s: Sample) -> Dict[str, Any]:
|
| assistant = {
|
| "criteria": {k: {"score": round(v, 1)} for k, v in s.scores.items()},
|
| "total_score": int(s.total_score),
|
| }
|
| return {
|
| "messages": [
|
| {"role": "user", "content": PROMPT_SCORE},
|
| {"role": "assistant", "content": json.dumps(assistant, ensure_ascii=False, indent=2)},
|
| ],
|
| "images": [s.image_path],
|
| }
|
|
|
|
|
| def build_multitask_record(s: Sample) -> Dict[str, Any]:
|
| assistant = {
|
| "criteria": {
|
| k: {"score": round(s.scores[k], 1), "level": s.levels[k]} for k in CRITERIA
|
| },
|
| "total_score": int(s.total_score),
|
| }
|
| return {
|
| "messages": [
|
| {"role": "user", "content": PROMPT_MULTITASK},
|
| {"role": "assistant", "content": json.dumps(assistant, ensure_ascii=False, indent=2)},
|
| ],
|
| "images": [s.image_path],
|
| }
|
|
|
|
|
| def compute_level_counts(samples: List[Sample]) -> Dict[str, Counter]:
|
| counts: Dict[str, Counter] = {c: Counter() for c in CRITERIA}
|
| for s in samples:
|
| for c in CRITERIA:
|
| counts[c][s.levels[c]] += 1
|
| return counts
|
|
|
|
|
| def compute_sample_dup(
|
| s: Sample,
|
| level_counts: Dict[str, Counter],
|
| max_class_weight: float,
|
| boundary_bonus: float,
|
| max_dup: int,
|
| ) -> int:
|
| weights: List[float] = []
|
| for c in CRITERIA:
|
| cnt = level_counts[c]
|
| max_freq = max(cnt.values()) if cnt else 1
|
| cur = cnt.get(s.levels[c], 1)
|
| w = math.sqrt(float(max_freq) / float(max(1, cur)))
|
| w = min(max_class_weight, max(1.0, w))
|
| weights.append(w)
|
|
|
| avg_w = sum(weights) / len(weights)
|
| raw_dup = avg_w + boundary_bonus * float(s.boundary_count)
|
| dup = int(round(raw_dup))
|
| return clamp_int(dup, 1, max_dup)
|
|
|
|
|
| def save_json(path: str, obj: Any) -> None:
|
| os.makedirs(os.path.dirname(path), exist_ok=True)
|
| with open(path, "w", encoding="utf-8") as f:
|
| json.dump(obj, f, ensure_ascii=False, indent=2)
|
|
|
|
|
| def build_stats(
|
| train_samples: List[Sample],
|
| val_samples: List[Sample],
|
| level_counts: Dict[str, Counter],
|
| dup_counts: Counter,
|
| balanced_size: int,
|
| ) -> Dict[str, Any]:
|
| agg = Counter()
|
| for c in CRITERIA:
|
| agg.update(level_counts[c])
|
|
|
| total_lv = sum(agg.values())
|
| agg_ratio = {k: round(v / total_lv, 6) for k, v in agg.items()} if total_lv else {}
|
|
|
| per_criterion = {}
|
| for c in CRITERIA:
|
| cnt = level_counts[c]
|
| n = sum(cnt.values())
|
| per_criterion[c] = {
|
| "counts": dict(cnt),
|
| "ratio": {k: round(cnt[k] / n, 6) if n else 0.0 for k in ["A", "B", "C"]},
|
| }
|
|
|
| train_total = [s.total_score for s in train_samples]
|
| val_total = [s.total_score for s in val_samples]
|
|
|
| return {
|
| "train_count": len(train_samples),
|
| "val_count": len(val_samples),
|
| "balanced_train_count": balanced_size,
|
| "level_distribution_aggregate": dict(agg),
|
| "level_distribution_aggregate_ratio": agg_ratio,
|
| "level_distribution_per_criterion": per_criterion,
|
| "duplication_histogram": {str(k): v for k, v in sorted(dup_counts.items())},
|
| "boundary_stats": {
|
| "train_mean_boundary_count": round(sum(s.boundary_count for s in train_samples) / max(1, len(train_samples)), 4),
|
| "train_max_boundary_count": max((s.boundary_count for s in train_samples), default=0),
|
| },
|
| "total_score": {
|
| "train_mean": round(sum(train_total) / max(1, len(train_total)), 4),
|
| "val_mean": round(sum(val_total) / max(1, len(val_total)), 4),
|
| "train_min": min(train_total) if train_total else None,
|
| "train_max": max(train_total) if train_total else None,
|
| "val_min": min(val_total) if val_total else None,
|
| "val_max": max(val_total) if val_total else None,
|
| },
|
| }
|
|
|
|
|
| def main() -> None:
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument(
|
| "--train_json",
|
| type=str,
|
| default="./codabench_portrait_train.json",
|
| )
|
| parser.add_argument(
|
| "--val_json",
|
| type=str,
|
| default="./codabench_portrait_val.json",
|
| )
|
| parser.add_argument("--output_dir", type=str, default=".")
|
| parser.add_argument("--seed", type=int, default=42)
|
| parser.add_argument("--boundary_margin", type=float, default=0.4)
|
| parser.add_argument("--boundary_bonus", type=float, default=0.3)
|
| parser.add_argument("--max_class_weight", type=float, default=4.0)
|
| parser.add_argument("--max_dup", type=int, default=5)
|
| parser.add_argument(
|
| "--stats_json",
|
| type=str,
|
| default="logs/prepare_dataset_v2_stats.json",
|
| )
|
| args = parser.parse_args()
|
|
|
| random.seed(args.seed)
|
|
|
| with open(args.train_json, "r", encoding="utf-8") as f:
|
| train_raw = json.load(f)
|
| with open(args.val_json, "r", encoding="utf-8") as f:
|
| val_raw = json.load(f)
|
|
|
| train_samples = [normalize_record(x, args.boundary_margin) for x in train_raw]
|
| val_samples = [normalize_record(x, args.boundary_margin) for x in val_raw]
|
|
|
| level_counts = compute_level_counts(train_samples)
|
|
|
| train_score = [build_score_record(s) for s in train_samples]
|
| train_multitask = [build_multitask_record(s) for s in train_samples]
|
| val_score = [build_score_record(s) for s in val_samples]
|
| val_multitask = [build_multitask_record(s) for s in val_samples]
|
|
|
| train_balanced: List[Dict[str, Any]] = []
|
| dup_hist = Counter()
|
| for s in train_samples:
|
| dup = compute_sample_dup(
|
| s=s,
|
| level_counts=level_counts,
|
| max_class_weight=args.max_class_weight,
|
| boundary_bonus=args.boundary_bonus,
|
| max_dup=args.max_dup,
|
| )
|
| dup_hist[dup] += 1
|
| rec = build_multitask_record(s)
|
| for _ in range(dup):
|
| train_balanced.append(rec)
|
|
|
| random.shuffle(train_balanced)
|
|
|
| os.makedirs(args.output_dir, exist_ok=True)
|
|
|
| f_train_score = os.path.join(args.output_dir, "codabench_portrait_train_score_v2.json")
|
| f_train_multitask = os.path.join(args.output_dir, "codabench_portrait_train_multitask_v2.json")
|
| f_train_balanced = os.path.join(args.output_dir, "codabench_portrait_train_balanced_v2.json")
|
| f_val_score = os.path.join(args.output_dir, "codabench_portrait_val_score_v2.json")
|
| f_val_multitask = os.path.join(args.output_dir, "codabench_portrait_val_multitask_v2.json")
|
|
|
| save_json(f_train_score, train_score)
|
| save_json(f_train_multitask, train_multitask)
|
| save_json(f_train_balanced, train_balanced)
|
| save_json(f_val_score, val_score)
|
| save_json(f_val_multitask, val_multitask)
|
|
|
| stats = build_stats(
|
| train_samples=train_samples,
|
| val_samples=val_samples,
|
| level_counts=level_counts,
|
| dup_counts=dup_hist,
|
| balanced_size=len(train_balanced),
|
| )
|
| save_json(args.stats_json, stats)
|
|
|
| print("=" * 60)
|
| print("saved:", f_train_score)
|
| print("saved:", f_train_multitask)
|
| print("saved:", f_train_balanced)
|
| print("saved:", f_val_score)
|
| print("saved:", f_val_multitask)
|
| print("stats:", args.stats_json)
|
| print("train_count:", len(train_samples), "balanced_count:", len(train_balanced), "val_count:", len(val_samples))
|
| print("dup_hist:", dict(sorted(dup_hist.items())))
|
| print("=" * 60)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
