0llheaven/Conditional_DETR_TF

1import gradio as gr2from transformers import AutoImageProcessor, AutoModelForObjectDetection3import torch4from PIL import Image, ImageDraw5 6# Load the model and processor7processor = AutoImageProcessor.from_pretrained("0llheaven/Conditional-detr-finetuned-tf")8model = AutoModelForObjectDetection.from_pretrained("0llheaven/Conditional-detr-finetuned-tf")9 10def detect_objects(image, score_threshold):11    # Convert image to RGB if it's grayscale12    if image.mode != "RGB":13        image = image.convert("RGB")14    15    # Prepare input for the model16    inputs = processor(images=image, return_tensors="pt")17    outputs = model(**inputs)18 19    # Filter predictions based on the user-defined score threshold20    target_sizes = torch.tensor([image.size[::-1]])21    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)22 23    labels_output = [] 24    25    # Draw bounding boxes around detected objects26    draw = ImageDraw.Draw(image)27    for result in results:28        scores = result["scores"]29        labels = result["labels"]30        boxes = result["boxes"]31 32        for score, label, box in zip(scores, labels, boxes):33            if score >= score_threshold:  # Only draw if score is above threshold34                box = [round(i, 2) for i in box.tolist()]35                label_name = "Pneumonia" if label.item() == 0 else "No detection"36                draw.rectangle(box, outline="red", width=3)37                draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")38                labels_output.append(f"{label_name}: {round(score.item(), 3)}")39 40    # If no objects detected, append "No detection"41    if not labels_output:42        labels_output.append("No detection")43    44    return image, "\n".join(labels_output)45 46# Create the Gradio interface47interface = gr.Interface(48    fn=detect_objects, 49    inputs=[gr.Image(type="pil"), gr.Slider(0, 1, value=0.5, label="Score Threshold")],  # Add slider for score threshold50    # outputs=gr.Image(type="pil"),  # Corrected output type51    outputs=[gr.Image(type="pil"), gr.Textbox(label="Detected Objects")],52    title="Object Detection with Transformers",53    description="Upload an image to detect objects using a fine-tuned Conditional-DETR model."54)55 56# Launch the interface57interface.launch()