167 lines
4.4 KiB
Python
167 lines
4.4 KiB
Python
# backend\ml\train_scene_model.py
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from torchvision import models, transforms
|
|
|
|
|
|
def load_encoder():
|
|
weights = models.ResNet18_Weights.DEFAULT
|
|
model = models.resnet18(weights=weights)
|
|
model.fc = torch.nn.Identity()
|
|
model.eval()
|
|
|
|
preprocess = weights.transforms()
|
|
return model, preprocess
|
|
|
|
|
|
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 {
|
|
"peopleCount": int(pred.get("peopleCount") or 0),
|
|
"maleCount": int(pred.get("maleCount") or 0),
|
|
"femaleCount": int(pred.get("femaleCount") or 0),
|
|
"unknownCount": int(pred.get("unknownCount") or 0),
|
|
"sexPosition": str(pred.get("sexPosition") or "unknown"),
|
|
"bodyPartsPresent": [x.get("label") for x in pred.get("bodyPartsPresent") or [] if x.get("label")],
|
|
"objectsPresent": [x.get("label") for x in pred.get("objectsPresent") or [] if x.get("label")],
|
|
"clothingPresent": [x.get("label") for x in pred.get("clothingPresent") or [] if x.get("label")],
|
|
}
|
|
|
|
|
|
def correction_target(annotation):
|
|
corr = annotation.get("correction") or {}
|
|
|
|
return {
|
|
"peopleCount": int(corr.get("peopleCount") or 0),
|
|
"maleCount": int(corr.get("maleCount") or 0),
|
|
"femaleCount": int(corr.get("femaleCount") or 0),
|
|
"unknownCount": int(corr.get("unknownCount") or 0),
|
|
"sexPosition": str(corr.get("sexPosition") or "unknown"),
|
|
"bodyPartsPresent": list(corr.get("bodyPartsPresent") or []),
|
|
"objectsPresent": list(corr.get("objectsPresent") or []),
|
|
"clothingPresent": list(corr.get("clothingPresent") or []),
|
|
}
|
|
|
|
|
|
def target_from_annotation(annotation):
|
|
if annotation.get("accepted") is True:
|
|
return prediction_target(annotation)
|
|
|
|
return correction_target(annotation)
|
|
|
|
|
|
def embed_image(model, preprocess, image_path: Path):
|
|
img = Image.open(image_path).convert("RGB")
|
|
x = preprocess(img).unsqueeze(0)
|
|
|
|
with torch.no_grad():
|
|
emb = model(x).cpu().numpy()[0].astype("float32")
|
|
|
|
norm = np.linalg.norm(emb)
|
|
if norm > 0:
|
|
emb = emb / norm
|
|
|
|
return emb
|
|
|
|
|
|
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)
|
|
|
|
model, preprocess = load_encoder()
|
|
|
|
embeddings = []
|
|
targets = []
|
|
used = 0
|
|
skipped = 0
|
|
|
|
for row in rows:
|
|
sample_id = str(row.get("sampleId") or "").strip()
|
|
if not sample_id:
|
|
skipped += 1
|
|
continue
|
|
|
|
image_path = frames_dir / f"{sample_id}.jpg"
|
|
if not image_path.exists():
|
|
skipped += 1
|
|
continue
|
|
|
|
try:
|
|
emb = embed_image(model, preprocess, image_path)
|
|
target = target_from_annotation(row)
|
|
|
|
embeddings.append(emb)
|
|
targets.append(target)
|
|
used += 1
|
|
except Exception:
|
|
skipped += 1
|
|
|
|
if used < 5:
|
|
raise SystemExit(f"too few usable samples: {used}")
|
|
|
|
embeddings_np = np.stack(embeddings).astype("float32")
|
|
|
|
np.savez_compressed(
|
|
model_dir / "scene_knn.npz",
|
|
embeddings=embeddings_np,
|
|
)
|
|
|
|
with (model_dir / "targets.json").open("w", encoding="utf-8") as f:
|
|
json.dump(targets, f, ensure_ascii=False, indent=2)
|
|
|
|
with (model_dir / "status.json").open("w", encoding="utf-8") as f:
|
|
json.dump(
|
|
{
|
|
"ok": True,
|
|
"usedSamples": used,
|
|
"skippedSamples": skipped,
|
|
"model": "resnet18_knn",
|
|
},
|
|
f,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
|
|
print(json.dumps({
|
|
"ok": True,
|
|
"usedSamples": used,
|
|
"skippedSamples": skipped,
|
|
"modelPath": str(model_dir / "scene_knn.npz"),
|
|
}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |