| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from threading import Lock |
| from pathlib import Path |
|
|
| from PIL import Image |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| from cropvlm import CROP_CLASSES, load_cropvlm |
| from cropvlm.model import parse_class_names |
|
|
|
|
| DEFAULT_CLASSES_TEXT = "\n".join(CROP_CLASSES) |
|
|
|
|
| def build_demo(checkpoint: str, device: str | None, prompt_template: str, top_k: int) -> gr.Blocks: |
| import gradio as gr |
|
|
| classifier = load_cropvlm( |
| checkpoint=checkpoint, |
| class_names=CROP_CLASSES, |
| device=device, |
| prompt_template=prompt_template, |
| ) |
| classifier_lock = Lock() |
| current_classes = tuple(CROP_CLASSES) |
|
|
| def classify(image: Image.Image, classes_text: str, top_k_value: int): |
| if image is None: |
| return {}, [] |
| nonlocal current_classes |
| requested_classes = tuple(parse_class_names(classes_text)) |
| if not requested_classes: |
| return {}, [] |
| with classifier_lock: |
| if requested_classes != current_classes: |
| classifier.set_classes(requested_classes) |
| current_classes = requested_classes |
| predictions = classifier.predict_with_scores(image, top_k=int(top_k_value)) |
|
|
| label_scores = {label: probability for label, probability, _ in predictions} |
| score_text = "\n".join( |
| f"{rank}. {label}: probability={probability:.6f}, cosine={cosine:.6f}" |
| for rank, (label, probability, cosine) in enumerate(predictions, start=1) |
| ) |
| return label_scores, score_text |
|
|
| examples_dir = Path(__file__).resolve().parents[1] / "examples" |
| example_paths = [ |
| str(examples_dir / name) |
| for name in ["cacao.png", "olive.png", "cauliflower.png", "sugarcane.png", "sunflower.png"] |
| if (examples_dir / name).exists() |
| ] |
|
|
| with gr.Blocks(title="CropVLM Zero-Shot Demo") as demo: |
| gr.Markdown("# CropVLM Zero-Shot Image Classification") |
| with gr.Row(): |
| with gr.Column(): |
| image = gr.Image(type="pil", label="Image") |
| classes = gr.Textbox( |
| value=DEFAULT_CLASSES_TEXT, |
| lines=12, |
| label="Class names", |
| ) |
| top_k_slider = gr.Slider( |
| minimum=1, |
| maximum=10, |
| value=top_k, |
| step=1, |
| label="Top-k", |
| ) |
| button = gr.Button("Classify", variant="primary") |
| with gr.Column(): |
| label = gr.Label(num_top_classes=top_k, label="Predictions") |
| score_text = gr.Textbox( |
| label="Scores", |
| lines=8, |
| interactive=False, |
| ) |
|
|
| outputs = [label, score_text] |
| button.click(classify, inputs=[image, classes, top_k_slider], outputs=outputs) |
| classes.change(lambda: ({}, ""), outputs=outputs) |
|
|
| if example_paths: |
| gr.Examples( |
| examples=[[path, DEFAULT_CLASSES_TEXT, top_k] for path in example_paths], |
| inputs=[image, classes, top_k_slider], |
| outputs=outputs, |
| fn=classify, |
| cache_examples=False, |
| ) |
|
|
| return demo |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--checkpoint", default="models/CropVLM.pth") |
| parser.add_argument("--device", default=None) |
| parser.add_argument("--prompt-template", default="{}") |
| parser.add_argument("--top-k", type=int, default=5) |
| parser.add_argument("--server-name", default="127.0.0.1") |
| parser.add_argument("--server-port", type=int, default=7860) |
| args = parser.parse_args() |
|
|
| demo = build_demo( |
| checkpoint=args.checkpoint, |
| device=args.device, |
| prompt_template=args.prompt_template, |
| top_k=args.top_k, |
| ) |
| demo.launch(server_name=args.server_name, server_port=args.server_port) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
