import torch
from torch import nn
import torch.nn.functional as F
from transformers import AutoProcessor, AutoTokenizer, XCLIPVisionModel, AutoModel, AutoModelForSequenceClassification
import numpy as np
import cv2
import opensmile
class TextClassificationModel:
def __init__(self, model, device):
self.model = model
self.device = device
self.model.to(device)
def __call__(self, input_ids, attn_mask, return_last_hidden_state=False):
self.model.eval()
with torch.no_grad():
input_ids = input_ids.to(self.device)
attn_mask = attn_mask.to(self.device)
output = self.model(input_ids=input_ids, attention_mask=attn_mask,
output_hidden_states=return_last_hidden_state)
logits = output['logits']
pred = torch.argmax(logits, dim=1)
if return_last_hidden_state:
hidden_states = output['hidden_states']
if return_last_hidden_state:
return pred, hidden_states[-1][:, 0, :]
else:
return pred
class XCLIPClassificationModel(nn.Module):
def __init__(self, num_labels):
super(XCLIPClassificationModel, self).__init__()
self.base_model = XCLIPVisionModel.from_pretrained("microsoft/xclip-base-patch32")
self.num_labels = num_labels
hidden_size = self.base_model.config.hidden_size
self.fc_norm = nn.LayerNorm(hidden_size)
self.classifier = nn.Linear(hidden_size, self.num_labels)
self.loss_fct = nn.CrossEntropyLoss()
self.pool1 = nn.AdaptiveAvgPool1d(1)
self.pool2 = nn.AdaptiveAvgPool1d(1)
def forward(self, pixel_values, labels=None, return_last_hidden_state=False):
batch_size, num_frames, num_channels, height, width = pixel_values.shape
pixel_values = pixel_values.reshape(-1, num_channels, height, width)
out = self.base_model(pixel_values)[0] # [48, 50, 768]
out = torch.transpose(out, 1, 2) # [48, 768, 50]
out = self.pool1(out) # [48, 768, 1]
out = torch.transpose(out, 1, 2) # [48, 1, 768]
out = out.squeeze(1) # [48, 768]
hidden_out = out.view(batch_size, num_frames, -1) # [3, 16, 768]
hidden_out = torch.transpose(hidden_out, 1, 2) # [3, 768, 16]
pooled_out = self.pool2(hidden_out) # [3, 768, 1]
pooled_out = torch.transpose(pooled_out, 1, 2) # [3, 1, 768]
pooled_out = pooled_out[:, 0, :] # [3, 768]
logits = self.classifier(pooled_out)
loss = None
if labels is not None:
loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
if return_last_hidden_state:
return {'logits': logits, 'loss': loss, 'last_hidden_state': pooled_out}
else:
return {'logits': logits, 'loss': loss}
class VideoClassificationModel:
def __init__(self, model, device):
self.model = model
self.device = device
self.model.to(device)
def __call__(self, pixel_values, return_last_hidden_state=False):
self.model.eval()
with torch.no_grad():
pixel_values = pixel_values.to(self.device)
output = self.model(pixel_values, return_last_hidden_state=return_last_hidden_state)
logits = output['logits']
pred = torch.argmax(logits, dim=1)
if return_last_hidden_state:
hidden_states = output['last_hidden_state']
if return_last_hidden_state:
return pred, hidden_states
else:
return pred
class ConvNet(nn.Module):
def __init__(self, num_labels, n_input=1, n_channel=32):
super(ConvNet, self).__init__()
self.ln0 = nn.LayerNorm((1, 6191))
self.conv1 = nn.Conv1d(n_input, n_channel, kernel_size=3)
self.conv2 = nn.Conv1d(n_channel, n_channel, kernel_size=3)
self.bn1 = nn.BatchNorm1d(n_channel)
self.bn2 = nn.BatchNorm1d(n_channel)
self.pool1 = nn.MaxPool1d(2)
self.fc1 = nn.Linear(n_channel*3093, 3093)
self.fc2 = nn.Linear(3093, num_labels)
self.flat = nn.Flatten()
self.dropout = nn.Dropout(0.3)
def forward(self, x, return_last_hidden_state=False):
x = self.ln0(x)
x = self.conv1(x)
x = F.relu(self.bn1(x))
x = self.conv2(x)
x = F.relu(self.bn2(x))
x = self.pool1(x)
x = self.dropout(x)
x = self.flat(x)
hid = F.relu(self.fc1(x))
x = self.fc2(hid)
if not return_last_hidden_state:
return {'logits': F.log_softmax(x, dim=1)}
else:
return {'logits': F.log_softmax(x, dim=1), 'last_hidden_state': hid}
class AudioClassificationModel:
def __init__(self, model, device):
self.model = model
self.device = device
self.model.to(device)
def __call__(self, input_ids, return_last_hidden_state=False):
self.model.eval()
with torch.no_grad():
input_ids = torch.tensor(input_ids, dtype=torch.float).to(self.device)
output = self.model(input_ids, return_last_hidden_state=return_last_hidden_state)
logits = output['logits']
pred = torch.argmax(logits, dim=1)
if return_last_hidden_state:
hidden_state = output['last_hidden_state']
if return_last_hidden_state:
return pred, hidden_state
else:
return pred
class MultimodalClassificationModel(nn.Module):
def __init__(self, text_model, video_model, audio_model, num_labels, input_size, hidden_size=256):
super(MultimodalClassificationModel, self).__init__()
self.text_model = text_model
self.video_model = video_model
self.audio_model = audio_model
self.num_labels = num_labels
self.linear1 = nn.Linear(input_size, hidden_size)
self.linear2 = nn.Linear(hidden_size, self.num_labels)
self.relu1 = nn.ReLU()
self.drop1 = nn.Dropout()
self.loss_func = nn.CrossEntropyLoss()
def forward(self, batch, labels=None):
text_pred, text_last_hidden = self.text_model(
batch['text']['input_ids'].squeeze(1),
batch['text']['attention_mask'].squeeze(1),
return_last_hidden_state=True
)
video_pred, video_last_hidden = self.video_model(
batch['video']['pixel_values'].squeeze(1),
return_last_hidden_state=True
)
audio_pred, audio_last_hidden = self.audio_model(
batch['audio'],
return_last_hidden_state=True
)
concat_input = torch.cat((text_last_hidden, video_last_hidden, audio_last_hidden), dim=1)
hidden_state = self.linear1(concat_input)
hidden_state = self.drop1(self.relu1(hidden_state))
logits = self.linear2(hidden_state)
loss = None
if labels is not None:
loss = self.loss_func(logits.view(-1, self.num_labels), labels.view(-1))
return {'logits': logits, 'loss': loss}
class MainModel:
def __init__(self, model, device):
self.model = model
self.device = device
self.model.to(device)
def __call__(self, batch):
self.model.eval()
with torch.no_grad():
output = self.model(batch)
logits = output['logits']
pred = torch.argmax(logits, dim=1)
return pred
def prepare_models(num_labels: int,
text_model_path: str,
video_model_path: str,
audio_model_path: str,
device: str='cpu'):
# TEXT
text_model_name = 'bert-large-uncased'
text_base_model = AutoModelForSequenceClassification.from_pretrained(
text_model_name,
num_labels=num_labels
)
state_dict = torch.load(text_model_path, map_location=torch.device('cpu'))
text_base_model.load_state_dict(state_dict, strict=False)
text_model = TextClassificationModel(text_base_model, device=device)
# VIDEO
video_base_model = XCLIPClassificationModel(num_labels)
state_dict = torch.load(video_model_path, map_location=torch.device('cpu'))
video_base_model.load_state_dict(state_dict, strict=False)
video_model = VideoClassificationModel(video_base_model, device=device)
# AUDIO
audio_base_model = ConvNet(num_labels)
checkpoint = torch.load(audio_model_path, map_location=torch.device('cpu'))
audio_base_model.load_state_dict(checkpoint['model_state_dict'])
audio_model = AudioClassificationModel(audio_base_model, device=device)
return text_model, video_model, audio_model
def sample_frame_indices(seg_len, clip_len=16, frame_sample_rate=4, mode="video"):
# seg_len -- how many frames are received
# clip_len -- how many frames to return
converted_len = int(clip_len * frame_sample_rate)
converted_len = min(converted_len, seg_len-1)
end_idx = np.random.randint(converted_len, seg_len)
start_idx = end_idx - converted_len
if mode == "video":
indices = np.linspace(start_idx, end_idx, num=clip_len)
else:
indices = np.linspace(start_idx, end_idx, num=clip_len*frame_sample_rate)
indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
return indices
def get_frames(file_path, clip_len=16,):
cap = cv2.VideoCapture(file_path)
v_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
indices = sample_frame_indices(v_len)
frames = []
for fn in range(v_len):
success, frame = cap.read()
if success is False:
continue
if (fn in indices):
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
res = cv2.resize(frame[90:-80, 60:-100], dsize=(224, 224), interpolation=cv2.INTER_CUBIC)
frames.append(res)
cap.release()
if len(frames) < clip_len:
add_num = clip_len - len(frames)
frames_to_add = [frames[-1]] * add_num
frames.extend(frames_to_add)
return frames
def prepare_data_input(text: str,
video_path: str):
# VIDEO
video_frames = get_frames(video_path)
video_model_name = "microsoft/xclip-base-patch32"
video_feature_extractor = AutoProcessor.from_pretrained(video_model_name)
video_encoding = video_feature_extractor(videos=video_frames, return_tensors="pt")
# AUDIO
smile = opensmile.Smile(
opensmile.FeatureSet.ComParE_2016,
opensmile.FeatureLevel.Functionals,
sampling_rate=16000,
resample=True,
num_workers=5,
verbose=True,
)
audio_features = smile.process_files([video_path])
redundant_feat = open('files/redundant_feat.txt').read().split(',')
audio_features.drop(columns=redundant_feat, inplace=True)
# TEXT
text_model_name = 'bert-large-uncased'
tokenizer = AutoTokenizer.from_pretrained(text_model_name)
text_encoding = tokenizer(text,
padding='max_length',
truncation=True,
max_length=128,
return_tensors='pt')
return {'text': text_encoding, 'video': video_encoding, 'audio': audio_features.values.reshape((1, 1, 6191))}
def infer_multimodal_model(text: str,
video_path: str,
model_pathes: dict):
label2id = {'anger': 0, 'disgust': 1, 'fear': 2, 'joy': 3, 'neutral': 4, 'sadness': 5, 'surprise': 6}
id2label = {v: k for k, v in label2id.items()}
num_labels = 7
text_model, video_model, audio_model = prepare_models(num_labels,
model_pathes['text_model_path'],
model_pathes['video_model_path'],
model_pathes['audio_model_path'],)
multi_model = MultimodalClassificationModel(
text_model,
video_model,
audio_model,
num_labels,
input_size=4885,
hidden_size=512
)
checkpoint = torch.load(model_pathes['multimodal_model_path'], map_location=torch.device('cpu'))
multi_model.load_state_dict(checkpoint)
device = 'cpu'
final_model = MainModel(multi_model, device=device)
batch = prepare_data_input(text, video_path)
label = final_model(batch).detach().cpu().tolist()
return id2label[label[0]]