Spaces:
Running
Running
# Transformers and its models | |
import transformers | |
# For Image Processing | |
from transformers import ViTImageProcessor | |
# For Model | |
from transformers import ViTModel, ViTConfig | |
# For data augmentation | |
from torchvision import transforms, datasets | |
# For GPU | |
from transformers import set_seed | |
from torch.optim import AdamW | |
from accelerate import Accelerator, notebook_launcher | |
# For Data Loaders | |
import datasets | |
from torch.utils.data import Dataset, DataLoader | |
# For Display | |
from tqdm.notebook import tqdm | |
# Other Generic Libraries | |
import torch | |
import PIL | |
import streamlit as st | |
import gc | |
from glob import glob | |
import shutil | |
import torch.nn.functional as F | |
# Initialse Globle Variables | |
MODEL_TRANSFORMER = 'google/vit-base-patch16-224' | |
BATCH_SIZE = 8 | |
# Set the device (GPU or CPU) | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog") | |
st.title("Hot Dog? Or Not?") | |
# Read images from directory | |
# Read image from Camera | |
enable = st.checkbox("Enable camera") | |
picture = st.camera_input("Take a picture", disabled=not enable) | |
if picture: | |
col1, col2 = st.columns(2) | |
image = PIL.Image.open(picture) | |
col1.image(image, use_column_width=True) | |
predictions = pipeline(image) | |
col2.header("Probabilities") | |
for p in predictions: | |
col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%") |