Views
No views yet
1import torch
2from transformers import AutoImageProcessor
3
4# Load
5checkpoint = torch.load('multihead_model.pt')
6processor = AutoImageProcessor.from_pretrained('path/to/model')
7
8# Create model class (see notebook for full class definition)
9# model = MultiHeadContentModerator(...)
10# model.load_state_dict(checkpoint['model_state_dict'])
11
12# Inference
13inputs = processor(images=image, return_tensors='pt')
14with torch.no_grad():
15 # Get both predictions
16 outputs = model(inputs['pixel_values'], task='both')
17 nsfw_pred = outputs['nsfw'].argmax(-1)
18 violence_pred = outputs['violence'].argmax(-1)