398 lines
12 KiB
Python
398 lines
12 KiB
Python
import argparse
|
|
import json
|
|
import math
|
|
import shutil
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from transformers import AutoImageProcessor, VideoMAEForVideoClassification
|
|
|
|
|
|
DEFAULT_BASE_MODEL = "MCG-NJU/videomae-base-finetuned-kinetics"
|
|
|
|
|
|
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 safe_int(value, fallback):
|
|
try:
|
|
return int(value)
|
|
except Exception:
|
|
return fallback
|
|
|
|
|
|
def safe_float(value, fallback):
|
|
try:
|
|
return float(value)
|
|
except Exception:
|
|
return fallback
|
|
|
|
|
|
def load_labels(dataset_root: Path) -> list[str]:
|
|
path = dataset_root / "labels.json"
|
|
if not path.exists():
|
|
raise SystemExit(f"labels.json not found: {path}")
|
|
|
|
with path.open("r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
labels = data.get("labels", [])
|
|
out = []
|
|
seen = set()
|
|
|
|
for value in labels:
|
|
label = str(value or "").strip()
|
|
if not label or label in seen:
|
|
continue
|
|
seen.add(label)
|
|
out.append(label)
|
|
|
|
if len(out) < 2:
|
|
raise SystemExit("VideoMAE needs at least two labels")
|
|
|
|
return out
|
|
|
|
|
|
def load_manifest(dataset_root: Path, label_to_id: dict[str, int]) -> list[dict]:
|
|
path = dataset_root / "manifest.jsonl"
|
|
if not path.exists():
|
|
raise SystemExit(f"manifest.jsonl not found: {path}")
|
|
|
|
entries = []
|
|
with path.open("r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
try:
|
|
item = json.loads(line)
|
|
except Exception:
|
|
continue
|
|
|
|
label = str(item.get("label") or "").strip()
|
|
clip_dir = Path(str(item.get("clipDir") or "")).expanduser()
|
|
split = str(item.get("split") or "").strip().lower()
|
|
|
|
if label not in label_to_id or split not in {"train", "val"}:
|
|
continue
|
|
if not clip_dir.exists() or not clip_dir.is_dir():
|
|
continue
|
|
|
|
frames = sorted(
|
|
p for p in clip_dir.iterdir()
|
|
if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
|
|
)
|
|
if not frames:
|
|
continue
|
|
|
|
item["label"] = label
|
|
item["split"] = split
|
|
item["clipDir"] = str(clip_dir)
|
|
item["frames"] = [str(p) for p in frames]
|
|
entries.append(item)
|
|
|
|
return entries
|
|
|
|
|
|
def resample_values(values: list, count: int) -> list:
|
|
if not values:
|
|
return []
|
|
if count <= 1:
|
|
return [values[0]]
|
|
if len(values) == 1:
|
|
return [values[0] for _ in range(count)]
|
|
|
|
out = []
|
|
last = len(values) - 1
|
|
for i in range(count):
|
|
idx = int(round((i * last) / max(1, count - 1)))
|
|
out.append(values[idx])
|
|
return out
|
|
|
|
|
|
def load_clip_frames(paths: list[str], num_frames: int) -> list[Image.Image]:
|
|
selected = resample_values(paths, num_frames)
|
|
frames = []
|
|
|
|
for path in selected:
|
|
with Image.open(path) as img:
|
|
frames.append(img.convert("RGB").copy())
|
|
|
|
return frames
|
|
|
|
|
|
class VideoMAEClipDataset(Dataset):
|
|
def __init__(self, entries, label_to_id, image_processor, num_frames):
|
|
self.entries = entries
|
|
self.label_to_id = label_to_id
|
|
self.image_processor = image_processor
|
|
self.num_frames = num_frames
|
|
|
|
def __len__(self):
|
|
return len(self.entries)
|
|
|
|
def __getitem__(self, idx):
|
|
item = self.entries[idx]
|
|
frames = load_clip_frames(item["frames"], self.num_frames)
|
|
inputs = self.image_processor(frames, return_tensors="pt")
|
|
pixel_values = inputs["pixel_values"].squeeze(0)
|
|
label_id = self.label_to_id[item["label"]]
|
|
return {
|
|
"pixel_values": pixel_values,
|
|
"labels": torch.tensor(label_id, dtype=torch.long),
|
|
"sample_id": str(item.get("sampleId") or ""),
|
|
}
|
|
|
|
|
|
def collate_batch(batch):
|
|
return {
|
|
"pixel_values": torch.stack([item["pixel_values"] for item in batch]),
|
|
"labels": torch.stack([item["labels"] for item in batch]),
|
|
"sample_ids": [item["sample_id"] for item in batch],
|
|
}
|
|
|
|
|
|
def evaluate(model, loader, device):
|
|
model.eval()
|
|
total = 0
|
|
correct = 0
|
|
loss_sum = 0.0
|
|
|
|
with torch.no_grad():
|
|
for batch in loader:
|
|
pixel_values = batch["pixel_values"].to(device)
|
|
labels = batch["labels"].to(device)
|
|
outputs = model(pixel_values=pixel_values, labels=labels)
|
|
logits = outputs.logits
|
|
loss = outputs.loss
|
|
|
|
preds = logits.argmax(dim=-1)
|
|
total += int(labels.numel())
|
|
correct += int((preds == labels).sum().item())
|
|
loss_sum += float(loss.item()) * int(labels.numel())
|
|
|
|
if total <= 0:
|
|
return 0.0, 0.0
|
|
|
|
return correct / total, loss_sum / total
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--base", default=DEFAULT_BASE_MODEL)
|
|
parser.add_argument("--epochs", default="8")
|
|
parser.add_argument("--batch-size", default="2")
|
|
parser.add_argument("--lr", default="5e-5")
|
|
parser.add_argument("--device", default="auto")
|
|
parser.add_argument("--workers", default="0")
|
|
parser.add_argument("--num-frames", default="16")
|
|
parser.add_argument("--freeze-backbone", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
root = Path(args.root).resolve()
|
|
dataset_root = root / "videomae" / "dataset"
|
|
out_dir = root / "videomae" / "model"
|
|
tmp_dir = root / "videomae" / "runs" / "model_tmp"
|
|
|
|
epochs = max(1, safe_int(args.epochs, 8))
|
|
batch_size = max(1, safe_int(args.batch_size, 2))
|
|
workers = max(0, safe_int(args.workers, 0))
|
|
num_frames = max(2, safe_int(args.num_frames, 16))
|
|
lr = max(1e-7, safe_float(args.lr, 5e-5))
|
|
|
|
if str(args.device).lower() == "auto":
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
else:
|
|
device = torch.device(args.device)
|
|
|
|
labels = load_labels(dataset_root)
|
|
label_to_id = {label: i for i, label in enumerate(labels)}
|
|
id_to_label = {i: label for label, i in label_to_id.items()}
|
|
entries = load_manifest(dataset_root, label_to_id)
|
|
train_entries = [item for item in entries if item["split"] == "train"]
|
|
val_entries = [item for item in entries if item["split"] == "val"]
|
|
|
|
emit_progress(
|
|
"videomae",
|
|
0.01,
|
|
"VideoMAE-Dataset wird geprueft...",
|
|
trainSamples=len(train_entries),
|
|
valSamples=len(val_entries),
|
|
epochs=epochs,
|
|
device=str(device),
|
|
)
|
|
|
|
if not train_entries:
|
|
raise SystemExit("no VideoMAE train clips found")
|
|
if not val_entries:
|
|
raise SystemExit("no VideoMAE val clips found")
|
|
|
|
emit_progress(
|
|
"videomae",
|
|
0.03,
|
|
"VideoMAE-Basismodell wird geladen...",
|
|
base=args.base,
|
|
labels=len(labels),
|
|
device=str(device),
|
|
)
|
|
|
|
image_processor = AutoImageProcessor.from_pretrained(args.base)
|
|
model = VideoMAEForVideoClassification.from_pretrained(
|
|
args.base,
|
|
num_labels=len(labels),
|
|
label2id=label_to_id,
|
|
id2label=id_to_label,
|
|
ignore_mismatched_sizes=True,
|
|
)
|
|
|
|
if args.freeze_backbone and hasattr(model, "videomae"):
|
|
for param in model.videomae.parameters():
|
|
param.requires_grad = False
|
|
|
|
model.to(device)
|
|
|
|
train_ds = VideoMAEClipDataset(train_entries, label_to_id, image_processor, num_frames)
|
|
val_ds = VideoMAEClipDataset(val_entries, label_to_id, image_processor, num_frames)
|
|
train_loader = DataLoader(
|
|
train_ds,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=workers,
|
|
collate_fn=collate_batch,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_ds,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=workers,
|
|
collate_fn=collate_batch,
|
|
)
|
|
|
|
optimizer = torch.optim.AdamW(
|
|
[p for p in model.parameters() if p.requires_grad],
|
|
lr=lr,
|
|
)
|
|
|
|
if tmp_dir.exists():
|
|
shutil.rmtree(tmp_dir)
|
|
tmp_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
best_accuracy = -1.0
|
|
best_loss = math.inf
|
|
best_epoch = 0
|
|
|
|
for epoch in range(1, epochs + 1):
|
|
model.train()
|
|
running_loss = 0.0
|
|
seen = 0
|
|
total_batches = max(1, len(train_loader))
|
|
|
|
for batch_idx, batch in enumerate(train_loader, start=1):
|
|
pixel_values = batch["pixel_values"].to(device)
|
|
labels_tensor = batch["labels"].to(device)
|
|
|
|
optimizer.zero_grad(set_to_none=True)
|
|
outputs = model(pixel_values=pixel_values, labels=labels_tensor)
|
|
loss = outputs.loss
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
batch_size_seen = int(labels_tensor.numel())
|
|
seen += batch_size_seen
|
|
running_loss += float(loss.item()) * batch_size_seen
|
|
|
|
completed = (epoch - 1) + min(1.0, batch_idx / total_batches)
|
|
emit_progress(
|
|
"videomae",
|
|
0.04 + 0.84 * (completed / max(1, epochs)),
|
|
"VideoMAE trainiert...",
|
|
epoch=epoch,
|
|
epochs=epochs,
|
|
sampleId=(batch["sample_ids"][0] if batch["sample_ids"] else ""),
|
|
trainSamples=len(train_entries),
|
|
valSamples=len(val_entries),
|
|
device=str(device),
|
|
loss=(running_loss / max(1, seen)),
|
|
)
|
|
|
|
val_accuracy, val_loss = evaluate(model, val_loader, device)
|
|
is_best = val_accuracy > best_accuracy or (
|
|
math.isclose(val_accuracy, best_accuracy) and val_loss < best_loss
|
|
)
|
|
|
|
if is_best:
|
|
best_accuracy = val_accuracy
|
|
best_loss = val_loss
|
|
best_epoch = epoch
|
|
model.save_pretrained(tmp_dir)
|
|
image_processor.save_pretrained(tmp_dir)
|
|
|
|
emit_progress(
|
|
"videomae",
|
|
0.88 + 0.08 * (epoch / max(1, epochs)),
|
|
"VideoMAE wird validiert...",
|
|
epoch=epoch,
|
|
epochs=epochs,
|
|
trainSamples=len(train_entries),
|
|
valSamples=len(val_entries),
|
|
device=str(device),
|
|
accuracy=val_accuracy,
|
|
loss=val_loss,
|
|
)
|
|
|
|
if best_epoch <= 0:
|
|
model.save_pretrained(tmp_dir)
|
|
image_processor.save_pretrained(tmp_dir)
|
|
best_epoch = epochs
|
|
|
|
status = {
|
|
"trainedAt": datetime.now(timezone.utc).isoformat(),
|
|
"epochs": epochs,
|
|
"bestEpoch": best_epoch,
|
|
"trainSamples": len(train_entries),
|
|
"valSamples": len(val_entries),
|
|
"numFrames": num_frames,
|
|
"baseModel": args.base,
|
|
"device": str(device),
|
|
"accuracy": best_accuracy if best_accuracy >= 0 else 0.0,
|
|
"loss": best_loss if math.isfinite(best_loss) else 0.0,
|
|
"labels": labels,
|
|
}
|
|
|
|
with (tmp_dir / "status.json").open("w", encoding="utf-8") as f:
|
|
json.dump(status, f, ensure_ascii=False, indent=2)
|
|
|
|
if out_dir.exists():
|
|
shutil.rmtree(out_dir)
|
|
shutil.copytree(tmp_dir, out_dir)
|
|
|
|
emit_progress(
|
|
"videomae",
|
|
1.0,
|
|
"VideoMAE-Training abgeschlossen.",
|
|
epoch=best_epoch,
|
|
epochs=epochs,
|
|
trainSamples=len(train_entries),
|
|
valSamples=len(val_entries),
|
|
device=str(device),
|
|
accuracy=status["accuracy"],
|
|
loss=status["loss"],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|