206 lines
8.9 KiB
Python
206 lines
8.9 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
import time
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
|
||
from PIL import Image, ImageDraw
|
||
|
||
from .. import db
|
||
from ..capabilities import get_capability_matrix
|
||
from ..config import settings
|
||
from ..jobs import create_job
|
||
from ..modules.dataset.service import create_dataset, dataset_dir, generate_yolo_dataset_yaml, validate_dataset
|
||
from ..modules.results.service import scan_results, scan_training_curves
|
||
from ..schemas import JobCreate
|
||
|
||
|
||
REPORT_PATH = settings.project_root / "var" / "agent_reports" / "user_agent_latest.json"
|
||
|
||
|
||
def _now_id() -> str:
|
||
return datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_%f")
|
||
|
||
|
||
def _polygon_line(class_id: int, points: list[tuple[float, float]], width: int, height: int) -> str:
|
||
normalized = []
|
||
for x, y in points:
|
||
normalized.extend([max(0, min(1, x / width)), max(0, min(1, y / height))])
|
||
return f"{class_id} " + " ".join(f"{value:.6f}" for value in normalized)
|
||
|
||
|
||
def _ellipse_points(cx: float, cy: float, rx: float, ry: float, count: int = 24) -> list[tuple[float, float]]:
|
||
return [
|
||
(cx + math.cos(index / count * math.tau) * rx, cy + math.sin(index / count * math.tau) * ry)
|
||
for index in range(count)
|
||
]
|
||
|
||
|
||
def _write_open_synthetic_dataset(dataset_name: str, count: int = 6) -> dict:
|
||
create_dataset(dataset_name, "CC0-style synthetic segmentation data generated by the user agent.")
|
||
root = dataset_dir(dataset_name)
|
||
width = 160
|
||
height = 128
|
||
samples = []
|
||
for index in range(count):
|
||
stem = f"open_shape_{index:02d}"
|
||
image = Image.new("RGB", (width, height), (20 + index * 8, 28, 34))
|
||
mask = Image.new("L", (width, height), 0)
|
||
overlay = Image.new("RGB", (width, height), (0, 0, 0))
|
||
draw = ImageDraw.Draw(image)
|
||
mask_draw = ImageDraw.Draw(mask)
|
||
overlay_draw = ImageDraw.Draw(overlay)
|
||
|
||
ellipse = _ellipse_points(54 + index * 7, 54, 24, 18)
|
||
rectangle = [(92, 70), (132, 70), (132, 104), (92, 104)]
|
||
draw.polygon(ellipse, fill=(108, 193, 112))
|
||
draw.polygon(rectangle, fill=(104, 168, 230))
|
||
mask_draw.polygon(ellipse, fill=1)
|
||
mask_draw.polygon(rectangle, fill=2)
|
||
overlay_draw.polygon(ellipse, fill=(108, 193, 112))
|
||
overlay_draw.polygon(rectangle, fill=(104, 168, 230))
|
||
|
||
image_path = root / "images" / f"{stem}.png"
|
||
mask_path = root / "masks" / f"{stem}.png"
|
||
label_path = root / "labels" / f"{stem}.txt"
|
||
image.save(image_path)
|
||
mask.save(mask_path)
|
||
label_path.write_text(
|
||
"\n".join(
|
||
[
|
||
_polygon_line(0, ellipse, width, height),
|
||
_polygon_line(1, rectangle, width, height),
|
||
]
|
||
)
|
||
+ "\n",
|
||
encoding="utf-8",
|
||
)
|
||
samples.append({"image": str(image_path), "mask": str(mask_path), "label": str(label_path)})
|
||
|
||
manifest = {
|
||
"dataset": dataset_name,
|
||
"license": "CC0 synthetic data generated locally by Seg Data Server Net user agent",
|
||
"classes": ["soft_organ", "instrument"],
|
||
"samples": samples,
|
||
}
|
||
(root / "open_synthetic_manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
return manifest
|
||
|
||
|
||
def _write_review_artifacts(dataset_name: str) -> dict:
|
||
root = dataset_dir(dataset_name)
|
||
output_root = settings.project_root / "var" / "custom_yolo_runs" / f"{dataset_name}_user_agent_review"
|
||
predict_dir = output_root / "predict"
|
||
heatmap_dir = output_root / "heatmap"
|
||
predict_dir.mkdir(parents=True, exist_ok=True)
|
||
heatmap_dir.mkdir(parents=True, exist_ok=True)
|
||
for image_path in sorted((root / "images").glob("*.png"))[:3]:
|
||
image = Image.open(image_path).convert("RGB")
|
||
mask = Image.open(root / "masks" / image_path.name).convert("L")
|
||
overlay = Image.blend(image, Image.merge("RGB", (mask.point(lambda p: p * 90), mask.point(lambda p: p * 50), mask.point(lambda p: p * 20))), 0.35)
|
||
overlay.save(predict_dir / f"{image_path.stem}_segmentation.png")
|
||
|
||
heat = Image.new("RGB", image.size, (0, 0, 40))
|
||
heat_draw = ImageDraw.Draw(heat)
|
||
heat_draw.ellipse((32, 24, 112, 96), fill=(255, 70, 30))
|
||
heat_draw.rectangle((82, 60, 150, 118), fill=(45, 220, 255))
|
||
Image.blend(image, heat, 0.5).save(heatmap_dir / f"{image_path.stem}_heatmap.png")
|
||
|
||
results_csv = output_root / "results.csv"
|
||
results_csv.write_text(
|
||
"\n".join(
|
||
[
|
||
"epoch,train/box_loss,train/seg_loss,metrics/mIoU",
|
||
"0,1.000,0.850,0.420",
|
||
"1,0.720,0.610,0.630",
|
||
"2,0.530,0.430,0.760",
|
||
]
|
||
)
|
||
+ "\n",
|
||
encoding="utf-8",
|
||
)
|
||
return {
|
||
"root": str(output_root),
|
||
"predict_dir": str(predict_dir),
|
||
"heatmap_dir": str(heatmap_dir),
|
||
"results_csv": str(results_csv),
|
||
}
|
||
|
||
|
||
def _wait_job(job_id: str, timeout_seconds: float = 10) -> dict | None:
|
||
deadline = time.time() + timeout_seconds
|
||
while time.time() < deadline:
|
||
job = db.get_job(job_id)
|
||
if job and job["status"] in {"success", "failed", "cancelled"}:
|
||
return job
|
||
time.sleep(0.2)
|
||
return db.get_job(job_id)
|
||
|
||
|
||
def _save_report(report: dict) -> dict:
|
||
REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||
REPORT_PATH.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
return report
|
||
|
||
|
||
def latest_user_agent_report() -> dict:
|
||
if not REPORT_PATH.exists():
|
||
return {"available": False, "agent": "user_agent", "passed": False}
|
||
return json.loads(REPORT_PATH.read_text(encoding="utf-8"))
|
||
|
||
|
||
def run_user_agent() -> dict:
|
||
"""Act like a first-time operator with a small open synthetic segmentation dataset."""
|
||
db.init_db()
|
||
run_id = _now_id()
|
||
dataset_name = f"user_agent_open_shapes_{run_id}"
|
||
data_manifest = _write_open_synthetic_dataset(dataset_name)
|
||
validation = validate_dataset(dataset_name)
|
||
yolo_yaml = generate_yolo_dataset_yaml(dataset_name, ["soft_organ", "instrument"])
|
||
artifacts = _write_review_artifacts(dataset_name)
|
||
mock_job = create_job(JobCreate(type="mock.echo", params={"message": f"user-agent checked {dataset_name}"}))
|
||
finished_job = _wait_job(mock_job["id"])
|
||
results = scan_results(limit=1000)
|
||
curves = scan_training_curves(limit=100)
|
||
capabilities = get_capability_matrix()
|
||
|
||
result_prefix = f"var/custom_yolo_runs/{dataset_name}_user_agent_review"
|
||
visible_artifacts = [item for item in results if item["relative_path"].startswith(result_prefix)]
|
||
visible_curves = [item for item in curves if item["relative_path"].startswith(result_prefix)]
|
||
checks = [
|
||
{"name": "synthetic_open_dataset_created", "passed": len(data_manifest["samples"]) >= 6},
|
||
{"name": "image_mask_pairs_ready", "passed": validation["ready"]["mask"] and validation["pairs"]["image_mask"] >= 6, "detail": validation["pairs"]},
|
||
{"name": "yolo_labels_ready", "passed": validation["ready"]["yolo"] and validation["counts"]["annotations"] >= 12, "detail": validation["counts"]},
|
||
{"name": "dataset_yaml_generated", "passed": Path(yolo_yaml["path"]).exists(), "detail": yolo_yaml["relative_path"]},
|
||
{"name": "job_runner_used", "passed": bool(finished_job and finished_job["status"] == "success"), "detail": finished_job},
|
||
{"name": "result_artifacts_visible", "passed": len(visible_artifacts) >= 4, "detail": [item["relative_path"] for item in visible_artifacts[:8]]},
|
||
{"name": "training_curve_visible", "passed": len(visible_curves) >= 1, "detail": [item["relative_path"] for item in visible_curves[:4]]},
|
||
{"name": "capability_matrix_still_ready", "passed": capabilities["passed"], "detail": capabilities["summary"]},
|
||
]
|
||
suggestions = [
|
||
"推理页已经能选择训练权重与数据集图片源；建议下一步加一个批量对比视图，把多个 best.pt 对同一图片的输出并排显示。",
|
||
"数据集页能发现 image/label/mask 配对问题；建议后续提供彩色 label 调色板在线编辑与一键灰度 mask 转换。",
|
||
"结果页能读取合成预测图、热度图和 loss CSV；建议为真实长训任务增加按 run_id 固定筛选的结果集合。",
|
||
]
|
||
report = {
|
||
"available": True,
|
||
"agent": "user_agent",
|
||
"passed": all(item["passed"] for item in checks),
|
||
"run_id": run_id,
|
||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||
"dataset": {
|
||
"name": dataset_name,
|
||
"root": validation["root"],
|
||
"license": data_manifest["license"],
|
||
"counts": validation["counts"],
|
||
"pairs": validation["pairs"],
|
||
"yaml": yolo_yaml["relative_path"],
|
||
},
|
||
"artifacts": artifacts,
|
||
"checks": checks,
|
||
"suggestions": suggestions,
|
||
}
|
||
return _save_report(report)
|