|
import streamlit as st |
|
|
|
from transformers import pipeline |
|
import torch |
|
from PIL import Image |
|
|
|
|
|
|
|
|
|
def load_clip(model_size='large'): |
|
if model_size == 'base': |
|
MODEL_name = 'openai/clip-vit-base-patch32' |
|
|
|
|
|
|
|
model = CLIPModel.from_pretrained(MODEL_name) |
|
processor = CLIPImageProcessor.from_pretrained(MODEL_name) |
|
|
|
return processor, model |
|
|
|
def inference_clip(options, image, processor, model): |
|
|
|
inputs = processor(text= options, images=image, return_tensors="pt", padding=True) |
|
with torch.no_grad(): |
|
outputs = model(**inputs) |
|
|
|
|
|
logits_per_image = outputs.logits_per_image |
|
probs = logits_per_image.softmax(dim=1) |
|
|
|
max_prob_idx = torch.argmax(probs) |
|
max_prob_option = options[max_prob_idx] |
|
max_prob = probs[max_prob_idx].item() |
|
return max_prob_option |
|
|
|
|
|
|
|
col_l, col_r = st.columns(2) |
|
|
|
|
|
model_name = "openai/clip-vit-large-patch14-336" |
|
classifier = pipeline("zero-shot-image-classification", model = model_name) |
|
|
|
|
|
|
|
|
|
with col_l: |
|
picture_file = st.file_uploader("Picture :", type=["jpg", "jpeg", "png"]) |
|
|
|
if picture_file is not None: |
|
image = Image.open(picture_file) |
|
st.image(image, caption='Please upload an image of the damage') |
|
|
|
|
|
with col_l: |
|
default_options = 'There is a car, There is no car' |
|
options = st.text_input(label="Please enter the classes", value=default_options).split(',') |
|
|
|
|
|
|
|
|
|
if st.button("Compute"): |
|
|
|
|
|
scores = classifier(image, |
|
candidate_labels = options) |
|
|
|
with col_r: |
|
|
|
st.dataframe(scores) |
|
|