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