327 lines
8.3 KiB
Python
327 lines
8.3 KiB
Python
# backend\ml\train_scene_model.py
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import joblib
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.neighbors import KNeighborsClassifier
|
|
from transformers import CLIPModel, CLIPProcessor
|
|
|
|
|
|
CLIP_MODEL_NAME = "openai/clip-vit-base-patch32"
|
|
|
|
|
|
def read_jsonl(path: Path):
|
|
if not path.exists():
|
|
return []
|
|
|
|
rows = []
|
|
with path.open("r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
try:
|
|
rows.append(json.loads(line))
|
|
except Exception:
|
|
pass
|
|
|
|
return rows
|
|
|
|
|
|
def prediction_target(annotation):
|
|
pred = annotation.get("prediction") or {}
|
|
return str(pred.get("sexPosition") or "unknown").strip() or "unknown"
|
|
|
|
|
|
def correction_target(annotation):
|
|
corr = annotation.get("correction") or {}
|
|
return str(corr.get("sexPosition") or "unknown").strip() or "unknown"
|
|
|
|
|
|
def target_from_annotation(annotation):
|
|
if annotation.get("accepted") is True:
|
|
return prediction_target(annotation)
|
|
|
|
return correction_target(annotation)
|
|
|
|
def emit_progress(stage, progress, message="", **extra):
|
|
out = {
|
|
"type": "progress",
|
|
"stage": stage,
|
|
"progress": max(0.0, min(1.0, float(progress))),
|
|
"message": message,
|
|
}
|
|
out.update(extra)
|
|
print(json.dumps(out, ensure_ascii=False), flush=True)
|
|
|
|
def load_clip():
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME)
|
|
model = CLIPModel.from_pretrained(CLIP_MODEL_NAME)
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
return model, processor, device
|
|
|
|
|
|
def image_features_to_tensor(model, out):
|
|
if torch.is_tensor(out):
|
|
return out
|
|
|
|
if hasattr(out, "image_embeds") and out.image_embeds is not None:
|
|
return out.image_embeds
|
|
|
|
if hasattr(out, "pooler_output") and out.pooler_output is not None:
|
|
emb = out.pooler_output
|
|
|
|
# Nur projizieren, wenn pooler_output noch die erwartete Eingangsgröße hat.
|
|
# Bei manchen Transformers-Versionen ist pooler_output bereits 512-dimensional.
|
|
projection = getattr(model, "visual_projection", None)
|
|
if projection is not None and hasattr(projection, "in_features"):
|
|
if emb.shape[-1] == projection.in_features:
|
|
emb = projection(emb)
|
|
|
|
return emb
|
|
|
|
if isinstance(out, (tuple, list)) and len(out) > 0:
|
|
first = out[0]
|
|
if torch.is_tensor(first):
|
|
return first
|
|
|
|
raise TypeError(f"Unsupported CLIP image feature output: {type(out)!r}")
|
|
|
|
|
|
def embed_image(model, processor, device, image_path: Path):
|
|
img = Image.open(image_path).convert("RGB")
|
|
|
|
inputs = processor(images=img, return_tensors="pt")
|
|
inputs = {k: v.to(device) for k, v in inputs.items()}
|
|
|
|
with torch.no_grad():
|
|
try:
|
|
out = model.get_image_features(**inputs)
|
|
except Exception:
|
|
out = model.vision_model(pixel_values=inputs["pixel_values"])
|
|
|
|
emb = image_features_to_tensor(model, out)
|
|
|
|
emb = emb.detach().cpu().numpy()[0].astype("float32")
|
|
|
|
norm = np.linalg.norm(emb)
|
|
if norm > 0:
|
|
emb = emb / norm
|
|
|
|
return emb
|
|
|
|
|
|
def train_lr_if_possible(x, y):
|
|
classes = sorted(set(y))
|
|
|
|
if len(classes) < 2:
|
|
return None
|
|
|
|
# Logistic Regression braucht mindestens zwei Klassen.
|
|
# class_weight hilft bei unausgeglichenen Positionen.
|
|
clf = LogisticRegression(
|
|
max_iter=2000,
|
|
class_weight="balanced",
|
|
solver="lbfgs",
|
|
)
|
|
|
|
clf.fit(x, y)
|
|
return clf
|
|
|
|
|
|
def train_knn(x, y):
|
|
n_neighbors = min(7, len(y))
|
|
|
|
clf = KNeighborsClassifier(
|
|
n_neighbors=n_neighbors,
|
|
metric="cosine",
|
|
weights="distance",
|
|
algorithm="brute",
|
|
)
|
|
|
|
clf.fit(x, y)
|
|
return clf
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
args = parser.parse_args()
|
|
|
|
root = Path(args.root)
|
|
feedback_path = root / "feedback.jsonl"
|
|
frames_dir = root / "frames"
|
|
model_dir = root / "model"
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
rows = read_jsonl(feedback_path)
|
|
total = max(1, len(rows))
|
|
|
|
emit_progress(
|
|
"scene",
|
|
0.02,
|
|
"CLIP-Modell wird geladen…",
|
|
totalSamples=len(rows),
|
|
)
|
|
|
|
clip_model, processor, device = load_clip()
|
|
|
|
emit_progress(
|
|
"scene",
|
|
0.08,
|
|
"CLIP-Embeddings werden vorbereitet…",
|
|
totalSamples=len(rows),
|
|
device=device,
|
|
)
|
|
|
|
embeddings = []
|
|
labels = []
|
|
targets = []
|
|
used = 0
|
|
skipped = 0
|
|
|
|
for idx, row in enumerate(rows, start=1):
|
|
sample_id = str(row.get("sampleId") or "").strip()
|
|
|
|
try:
|
|
if not sample_id:
|
|
skipped += 1
|
|
continue
|
|
|
|
image_path = frames_dir / f"{sample_id}.jpg"
|
|
if not image_path.exists():
|
|
skipped += 1
|
|
continue
|
|
|
|
label = target_from_annotation(row)
|
|
if not label or label == "unknown":
|
|
skipped += 1
|
|
continue
|
|
|
|
emb = embed_image(clip_model, processor, device, image_path)
|
|
|
|
embeddings.append(emb)
|
|
labels.append(label)
|
|
targets.append({
|
|
"sampleId": sample_id,
|
|
"sexPosition": label,
|
|
})
|
|
used += 1
|
|
|
|
except Exception as e:
|
|
print(f"skip {sample_id or '<missing>'}: {repr(e)}", flush=True)
|
|
skipped += 1
|
|
|
|
finally:
|
|
emit_progress(
|
|
"scene",
|
|
0.08 + 0.78 * (idx / total),
|
|
f"Scene-Samples werden verarbeitet… {idx}/{len(rows)}",
|
|
currentSample=idx,
|
|
totalSamples=len(rows),
|
|
usedSamples=used,
|
|
skippedSamples=skipped,
|
|
)
|
|
|
|
if used < 5:
|
|
raise SystemExit(f"too few usable samples: {used}")
|
|
|
|
emit_progress(
|
|
"scene",
|
|
0.88,
|
|
"Scene-Embeddings werden gespeichert…",
|
|
usedSamples=used,
|
|
skippedSamples=skipped,
|
|
)
|
|
|
|
x = np.stack(embeddings).astype("float32")
|
|
y = np.array(labels)
|
|
|
|
np.savez_compressed(
|
|
model_dir / "scene_clip_embeddings.npz",
|
|
embeddings=x,
|
|
labels=y,
|
|
)
|
|
|
|
with (model_dir / "scene_clip_targets.json").open("w", encoding="utf-8") as f:
|
|
json.dump(targets, f, ensure_ascii=False, indent=2)
|
|
|
|
emit_progress(
|
|
"scene",
|
|
0.93,
|
|
"Scene-KNN wird trainiert…",
|
|
usedSamples=used,
|
|
skippedSamples=skipped,
|
|
)
|
|
|
|
knn = train_knn(x, y)
|
|
joblib.dump(knn, model_dir / "scene_clip_knn.joblib")
|
|
|
|
emit_progress(
|
|
"scene",
|
|
0.96,
|
|
"Scene-Logistic-Regression wird trainiert…",
|
|
usedSamples=used,
|
|
skippedSamples=skipped,
|
|
)
|
|
|
|
lr_status = "skipped_single_class"
|
|
lr = train_lr_if_possible(x, y)
|
|
if lr is not None:
|
|
joblib.dump(lr, model_dir / "scene_clip_lr.joblib")
|
|
lr_status = "trained"
|
|
else:
|
|
old_lr = model_dir / "scene_clip_lr.joblib"
|
|
if old_lr.exists():
|
|
old_lr.unlink()
|
|
|
|
counts = {}
|
|
for label in labels:
|
|
counts[label] = counts.get(label, 0) + 1
|
|
|
|
status = {
|
|
"ok": True,
|
|
"usedSamples": used,
|
|
"skippedSamples": skipped,
|
|
"model": "scene_position_clip",
|
|
"clipModel": CLIP_MODEL_NAME,
|
|
"device": device,
|
|
"classes": sorted(counts.keys()),
|
|
"classCounts": counts,
|
|
"logisticRegression": lr_status,
|
|
"knn": "trained",
|
|
"embeddingsPath": str(model_dir / "scene_clip_embeddings.npz"),
|
|
"knnPath": str(model_dir / "scene_clip_knn.joblib"),
|
|
"lrPath": str(model_dir / "scene_clip_lr.joblib"),
|
|
}
|
|
|
|
with (model_dir / "status.json").open("w", encoding="utf-8") as f:
|
|
json.dump(status, f, ensure_ascii=False, indent=2)
|
|
|
|
emit_progress(
|
|
"scene",
|
|
1.0,
|
|
"CLIP-Scene-Positionsmodell fertig.",
|
|
usedSamples=used,
|
|
skippedSamples=skipped,
|
|
classes=sorted(counts.keys()),
|
|
logisticRegression=lr_status,
|
|
knn="trained",
|
|
)
|
|
|
|
print(json.dumps(status, ensure_ascii=False), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |