Spaces:
Sleeping
Sleeping
File size: 5,651 Bytes
10fc801 26e8604 74905bc 26e8604 74905bc 10fc801 26e8604 10fc801 26e8604 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import gradio as gr
from transformers import AutoModel, AutoProcessor
from PIL import Image, ImageDraw, ImageFont
import logging
from datasets import load_dataset
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class DFineDemo:
def __init__(self):
self.processor = AutoProcessor.from_pretrained("Laudando-Associates-LLC/d-fine", trust_remote_code=True)
self.model_variants = {
"D-FINE Nano": "Laudando-Associates-LLC/d-fine-nano",
"D-FINE Small": "Laudando-Associates-LLC/d-fine-small",
"D-FINE Medium": "Laudando-Associates-LLC/d-fine-medium",
"D-FINE Large": "Laudando-Associates-LLC/d-fine-large",
"D-FINE X-Large": "Laudando-Associates-LLC/d-fine-xlarge"
}
logger.info("Loading all D-FINE model variants into memory...")
self.models = {
name: AutoModel.from_pretrained(repo, trust_remote_code=True)
for name, repo in self.model_variants.items()
}
dataset = load_dataset("Laudando-Associates-LLC/pucks", split="test")
self.image_cache = {
f"Test Image {i+1}": {
"input": example["image"],
"annotated": example["annotated_image"]
}
for i, example in enumerate(dataset)
}
self.image_labels = list(self.image_cache.keys())
def run_inference(self, input_image, model_name, threshold):
# Find matching annotated image based on value in self.image_cache
for label, pair in self.image_cache.items():
if pair["input"] == input_image:
annotated = pair["annotated"]
break
else:
annotated = input_image # fallback
# Predict
image = input_image.copy()
inputs = self.processor(image)
outputs = self.models[model_name](**inputs, conf_threshold=threshold)
draw = ImageDraw.Draw(image)
font = ImageFont.truetype("DejaVuSans-Bold.ttf", size=24)
for result in outputs:
for box, score in zip(result["boxes"], result["scores"]):
x1, y1, x2, y2 = box.tolist()
draw.rectangle([x1, y1, x2, y2], outline="blue", width=5)
draw.text((x1, max(0, y1 - 25)), f"{score:.2f}", fill="blue", font=font)
# Return: (annotated_image, predicted_image)
return gr.update(value=(annotated, image), slider_position=50, format="png", type="pil")
def select_image(self, evt: gr.SelectData):
if evt is None or evt.index is None:
return gr.update()
label = self.image_labels[evt.index]
return self.image_cache[label]["input"]
def launch(self):
with gr.Blocks(theme=gr.themes.Ocean()) as demo:
gr.Markdown("""
## D-FINE Detection Demo
This demo compares annotated ground truth data (in **red**) and model predictions (in **blue**).
Use the **slider** to visually compare both views:
- The **left image** shows the annotated labels.
- The **right image** displays predictions from the selected D-FINE model, with each bounding box and its confidence score.
📂 **Training Dataset**: All D-FINE variants were trained on the [L&A Pucks Dataset](https://huggingface.co/datasets/Laudando-Associates-LLC/pucks) available on Hugging Face.
""")
output = gr.ImageSlider(type="pil", label="Detected Output", height=500, width=880, slider_position=50, format="png")
with gr.Row():
model_selector = gr.Radio(
choices=list(self.model_variants.keys()),
label="Choose D-FINE model",
value="D-FINE Nano"
)
threshold_slider = gr.Slider(
minimum=0.1,
maximum=0.955,
value=0.4,
step=0.05,
label="Confidence Threshold"
)
run_btn = gr.Button("Run Detection")
selected_image = gr.State(value=self.image_cache[self.image_labels[0]])
gr.Markdown("### Select a sample image below:")
gallery = gr.Gallery(
value=[(pair["input"], label) for label, pair in self.image_cache.items()],
label=None,
show_label=False,
columns=[3],
object_fit="cover",
height="auto",
allow_preview=False
)
gallery.select(
fn=self.select_image,
inputs=[],
outputs=selected_image
)
run_btn.click(
fn=self.run_inference,
inputs=[selected_image, model_selector, threshold_slider],
outputs=output
)
gr.Markdown("### Citation")
gr.Markdown("""
If you use **D-FINE** or its methods in your work, please cite the following BibTeX entry:
```latex
@misc{peng2024dfine,
title={D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement},
author={Yansong Peng and Hebei Li and Peixi Wu and Yueyi Zhang and Xiaoyan Sun and Feng Wu},
year={2024},
eprint={2410.13842},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
```
""")
demo.launch()
if __name__ == "__main__":
app = DFineDemo()
app.launch()
|