|
|
from transformers import ViTImageProcessor |
|
|
import torch |
|
|
from PIL import Image |
|
|
import numpy as np |
|
|
|
|
|
def load_image_vit(image_file, processor): |
|
|
image = Image.open(image_file).convert('RGB') |
|
|
inputs = processor(images=image, return_tensors="pt") |
|
|
return inputs["pixel_values"] |
|
|
|
|
|
def predict_toxicity_vit(model, inputs, device): |
|
|
model.eval() |
|
|
with torch.no_grad(): |
|
|
inputs = inputs.to(device) |
|
|
outputs = model(inputs).logits |
|
|
probabilities = torch.nn.functional.softmax(outputs, dim=1) |
|
|
prediction = torch.argmax(probabilities, dim=1) |
|
|
return prediction.item(), probabilities[0].cpu().numpy() |
|
|
|
|
|
def get_label(prediction): |
|
|
return "Toxic" if prediction == 1 else "Non-Toxic" |
