|
--- |
|
license: apache-2.0 |
|
--- |
|
An vit classifier for handling noise image like this |
|
|
|
 |
|
|
|
It has limitation inbetween clear and noise |
|
|
|
``` |
|
from datasets import load_dataset |
|
from PIL import Image |
|
from transformers import ViTImageProcessor, ViTForImageClassification, TrainingArguments, Trainer |
|
import torch |
|
import numpy as np |
|
from datasets import load_metric |
|
import os |
|
import shutil |
|
|
|
model_name_or_path = 'lrzjason/noise-classifier' |
|
image_processor = ViTImageProcessor.from_pretrained(model_name_or_path) |
|
model = ViTForImageClassification.from_pretrained(model_name_or_path) |
|
|
|
input_dir = '' |
|
file = 'b5b457f4-5b52-4d68-be1b-9a2f557465f6.jpg' |
|
image = Image.open(os.path.join(input_dir, file)) |
|
|
|
inputs = image_processor(image, return_tensors="pt") |
|
with torch.no_grad(): |
|
logits = model(**inputs).logits |
|
|
|
# model predicts one of the 1000 ImageNet classes |
|
predicted_label = logits.argmax(-1).item() |
|
``` |
|
|