| 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) |
|
|