File size: 619 Bytes
66cb296
 
da54acc
66cb296
da54acc
 
 
66cb296
da54acc
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import gradio as gr

from transformers import pipeline, AutoModelForImageClassification, AutoFeatureExtractor

HF_MODEL_PATH = (
    "ImageIN/convnext-base-224_finetuned_on_unlabelled_IA_with_snorkel_labels"
)

classif_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_PATH)
feature_extractor = AutoFeatureExtractor.from_pretrained(HF_MODEL_PATH)

classif_pipeline = pipeline(
    "image-classification", model=classif_model, feature_extractor=feature_extractor
)

demo = gr.Interface(
    fn=lambda x: classif_pipeline(x)[0]["label"],
    inputs=gr.Image(type="pil"),
    outputs="text",
)
demo.launch()