|
|
|
|
|
import argparse, json, os |
|
|
import numpy as np |
|
|
import torch |
|
|
from huggingface_hub import snapshot_download |
|
|
from transformers import VideoMAEImageProcessor, AutoModel |
|
|
from decord import VideoReader |
|
|
|
|
|
ID2LABEL = {0: "normal", 1: "abnormal"} |
|
|
|
|
|
class ClassificationHead(torch.nn.Module): |
|
|
def __init__(self, in_dim: int, hidden_dims, num_labels: int = 2, dropout: float = 0.1): |
|
|
super().__init__() |
|
|
dims = [in_dim] + list(hidden_dims) |
|
|
layers = [] |
|
|
for i in range(len(dims) - 1): |
|
|
layers.append(torch.nn.Linear(dims[i], dims[i + 1])) |
|
|
layers.append(torch.nn.GELU()) |
|
|
layers.append(torch.nn.Dropout(dropout)) |
|
|
layers.append(torch.nn.Linear(dims[-1], num_labels)) |
|
|
self.net = torch.nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.net(x) |
|
|
|
|
|
def pool_tokens(feats: torch.Tensor, expected_feat_dim: int | None = None) -> torch.Tensor: |
|
|
if feats.dim() != 3: |
|
|
return feats |
|
|
_, d1, d2 = feats.shape |
|
|
if expected_feat_dim is not None: |
|
|
if d1 == expected_feat_dim: |
|
|
return feats.mean(dim=2) |
|
|
if d2 == expected_feat_dim: |
|
|
return feats.mean(dim=1) |
|
|
return feats.mean(dim=2 if d1 <= d2 else 1) |
|
|
|
|
|
def load_frames(path: str, clip_len: int): |
|
|
vr = VideoReader(path) |
|
|
idxs = np.linspace(0, len(vr) - 1, num=clip_len, dtype=int) |
|
|
return [vr[i].asnumpy() for i in idxs] |
|
|
|
|
|
def main(): |
|
|
ap = argparse.ArgumentParser() |
|
|
ap.add_argument("--repo_id", default="happy8825/internvideo_tuned") |
|
|
ap.add_argument("--video", required=True, help="Path to video file") |
|
|
ap.add_argument("--device", default="cuda", help="cuda or cpu") |
|
|
ap.add_argument("--precision", choices=["fp16","bf16","fp32"], default="bf16") |
|
|
args = ap.parse_args() |
|
|
|
|
|
device = torch.device(args.device if torch.cuda.is_available() else "cpu") |
|
|
local = snapshot_download(args.repo_id) |
|
|
cfg = json.load(open(os.path.join(local, "train_config.json"), "r")) |
|
|
base_model = cfg.get("base_model", "revliter/internvideo_next_large_p14_res224_f16") |
|
|
clip_len = int(cfg.get("clip_len", 16)) |
|
|
frame_size = int(cfg.get("frame_size", 224)) |
|
|
hidden = cfg.get("hidden", [512]) |
|
|
feature_dim = cfg.get("feature_dim") or cfg.get("hidden_size") |
|
|
|
|
|
processor = VideoMAEImageProcessor.from_pretrained(base_model) |
|
|
backbone = AutoModel.from_pretrained(base_model, trust_remote_code=True).to(device).eval() |
|
|
if device.type == "cuda" and args.precision != "fp32": |
|
|
target = torch.float16 if args.precision == "fp16" else torch.bfloat16 |
|
|
backbone.to(dtype=target) |
|
|
head = ClassificationHead(in_dim=feature_dim or backbone.config.hidden_size, hidden_dims=hidden) |
|
|
state = torch.load(os.path.join(local, "best_head.pt"), map_location="cpu") |
|
|
head.load_state_dict(state["head"]) |
|
|
head.to(device).eval() |
|
|
|
|
|
frames = load_frames(args.video, clip_len=clip_len) |
|
|
px = processor(frames, return_tensors="pt")["pixel_values"].permute(0,2,1,3,4).to(device) |
|
|
amp_dtype = torch.float16 if args.precision == "fp16" else (torch.bfloat16 if args.precision == "bf16" else None) |
|
|
with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=amp_dtype, enabled=(device.type=="cuda" and amp_dtype is not None)): |
|
|
feats = backbone.extract_features(pixel_values=px) |
|
|
pooled = pool_tokens(feats, expected_feat_dim=feature_dim) |
|
|
logits = head(pooled.float()) |
|
|
pred_id = int(logits.argmax(dim=-1).item()) |
|
|
print(ID2LABEL.get(pred_id, str(pred_id))) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
