Spaces:
Running
on
Zero
Running
on
Zero
| from transformers import VideoMAEFeatureExtractor, VideoMAEForVideoClassification | |
| import torch | |
| import numpy as np | |
| class ActPredScorer(torch.nn.Module): | |
| def __init__(self, model_name = "MCG-NJU/videomae-base-finetuned-kinetics", num_frames = 16, device = 'cuda', dtype=torch.float32): | |
| super().__init__() | |
| self.model = VideoMAEForVideoClassification.from_pretrained(model_name, num_frames = num_frames, torch_dtype=dtype) | |
| self.feature_extractor = VideoMAEFeatureExtractor.from_pretrained(model_name) | |
| self.device = device | |
| self.model.to(device) | |
| def get_target_class_idx(self, target_action): | |
| def mapping_func(x): | |
| if 'piano' in x: | |
| return 'playing piano' | |
| if 'guitar' in x: | |
| return 'playing guitar' | |
| if 'doughnuts' in x: | |
| return 'eating doughnuts' | |
| if 'beer' in x: | |
| return 'drinking beer' | |
| if 'badminton' in x: | |
| return 'playing badminton' | |
| if 'cello' in x: | |
| return 'playing cello' | |
| if 'scooter' in x: | |
| return 'riding scooter' | |
| if 'ballet' in x: | |
| return 'dancing ballet' | |
| if 'pancake' in x: | |
| return 'flipping pancake' | |
| if 'violin' in x: | |
| return 'playing violin' | |
| if 'wood' in x: | |
| return 'chopping wood' | |
| if 'watermelon' in x: | |
| return 'eating watermelon' | |
| if 'jogging' in x: | |
| return 'jogging' | |
| else: | |
| print(f"Please add your action mapping to ActPredScorer. Mapping not found for {x}") | |
| raise NotImplementedError | |
| try: | |
| target_class_idx = self.model.config.label2id[target_action] | |
| except: | |
| target_class_idx = self.model.config.label2id[mapping_func(target_action)] | |
| return target_class_idx | |
| def get_loss_and_score(self, norm_vid, target_action): | |
| ''' video should be a torch array of dtype float, with values from 0-1, of dimension (num_frames, height, width, 3)''' | |
| target_class_idx = self.get_target_class_idx(target_action) | |
| outputs = self.model(norm_vid, labels = torch.tensor([target_class_idx]).to(self.device)) | |
| loss = outputs.loss | |
| logits = outputs.logits | |
| norm_logits = torch.exp(logits)/ (torch.exp(logits).sum()) | |
| norm_logits = norm_logits.squeeze() | |
| score = norm_logits[target_class_idx] | |
| return loss, score, self.get_pred_class(logits) | |
| def get_pred_class(self, logits): | |
| predicted_class_idx = logits.argmax(-1).item() | |
| return self.model.config.id2label[predicted_class_idx] | |
| def gen_rand_labels_file(labels_list, out_file, num_labels = 50): | |
| idxs = np.random.choice(len(labels_list), num_labels, replace = False) | |
| rand_labels = [labels_list[i] for i in idxs] | |
| rand_labels.sort() | |
| with open(out_file, 'w') as f: | |
| for line in rand_labels: | |
| f.write(f"{line}\n") | |
| if __name__ == '__main__': | |
| # import numpy as np | |
| # scorer = ActPredScorer(num_frames = 7) | |
| # video_torch = [torch.randn((3,256,256)).clamp(0,1) for _ in range(7)] | |
| # encoding = scorer.feature_extractor(video_torch, do_rescale = False, return_tensors="pt") | |
| # print(scorer.get_loss_and_score(video_torch)) | |
| scorer = ActPredScorer(num_frames = 7) | |
| labels = scorer.model.config.id2label | |