1from transformers import AutoImageProcessor, ViTForImageClassification
2import torch
3from PIL import Image
4from pillow_heif import register_heif_opener, register_avif_opener
5
6register_heif_opener()
7register_avif_opener()
8
9def get_prediction(img):
10 image = Image.open(img).convert('RGB')
11 image_processor = AutoImageProcessor.from_pretrained("AashishKumar/AIvisionGuard-v2")
12 model = ViTForImageClassification.from_pretrained("AashishKumar/AIvisionGuard-v2")
13 inputs = image_processor(image, return_tensors="pt")
14
15 with torch.no_grad():
16 logits = model(**inputs).logits
17
18 top2_labels = logits.topk(2).indices.squeeze().tolist()
19 top2_scores = logits.topk(2).values.squeeze().tolist()
20
21 response = [{"label": model.config.id2label[label], "score": score} for label, score in zip(top2_labels, top2_scores)]
22 return response
The model was evaluated using the CIFake test dataset, with the following metrics:
This model provides a highly effective tool for detecting AI-generated synthetic images and has promising applications in content moderation, digital forensics, and trust preservation. Future improvements may include: