| import gradio as gr |
| import requests |
| import io |
| import os |
| from PIL import Image, ImageDraw, ImageFont |
| from pathlib import Path |
|
|
| API_URL = os.getenv("API_URL") |
| API_KEY = os.getenv("API_KEY") |
| IMAGE_FOLDER = os.getenv("IMAGE_FOLDER", "images") |
|
|
| def get_test_images(): |
| images = [] |
| if os.path.exists(IMAGE_FOLDER): |
| for file in sorted(Path(IMAGE_FOLDER).glob("*")): |
| if file.suffix.lower() in [".jpg", ".jpeg", ".png", ".bmp", ".gif"]: |
| images.append((str(file), file.name)) |
| return images |
|
|
| def load_test_image(image_path): |
| if image_path and os.path.exists(image_path): |
| return Image.open(image_path) |
| return None |
|
|
| CLASS_COLORS = { |
| 0: (255, 165, 0), |
| 1: (255, 0, 0), |
| 2: (0, 255, 0), |
| 3: (0, 128, 255), |
| 4: (255, 0, 255), |
| 5: (0, 0, 255), |
| 6: (255, 255, 0), |
| 7: (0, 255, 255), |
| } |
|
|
| def draw_boxes_on_image(image, detections): |
| if not detections: |
| return image |
|
|
| img_copy = image.copy() |
| draw = ImageDraw.Draw(img_copy) |
| img_width, img_height = img_copy.size |
|
|
| try: |
| font = ImageFont.truetype("arial.ttf", 20) |
| except: |
| font = ImageFont.load_default() |
|
|
| for detection in detections: |
| name = detection.get("name", "Unknown") |
| confidence = detection.get("confidence", 0) |
| class_id = detection.get("class", 0) |
| box = detection.get("box", {}) |
|
|
| color = CLASS_COLORS.get(class_id, (255, 255, 255)) |
|
|
| points = [ |
| (box.get("x1", 0), box.get("y1", 0)), |
| (box.get("x2", 0), box.get("y2", 0)), |
| (box.get("x3", 0), box.get("y3", 0)), |
| (box.get("x4", 0), box.get("y4", 0)) |
| ] |
|
|
| draw.polygon(points, outline=color, width=2) |
|
|
| if points: |
| y_coords = [p[1] for p in points] |
| x_coords = [p[0] for p in points] |
| max_y = max(y_coords) |
| min_x = min(x_coords) |
| max_x = max(x_coords) |
| center_x = (min_x + max_x) / 2 |
|
|
| label = f"{name} {confidence:.2f}" |
|
|
| bbox = draw.textbbox((0, 0), label, font=font) |
| text_width = bbox[2] - bbox[0] |
| text_height = bbox[3] - bbox[1] |
|
|
| label_x = center_x - text_width / 2 |
| label_y = max_y + 3 |
|
|
| if label_x < 0: |
| label_x = 2 |
| if label_x + text_width > img_width: |
| label_x = img_width - text_width - 2 |
| if label_y + text_height > img_height: |
| label_y = max_y - text_height - 3 |
|
|
| bg_padding = 2 |
| bg_box = [ |
| label_x - bg_padding, |
| label_y - bg_padding, |
| label_x + text_width + bg_padding, |
| label_y + text_height + bg_padding |
| ] |
| draw.rectangle(bg_box, outline=color, fill=(0, 0, 0)) |
|
|
| draw.text((label_x, label_y), label, font=font) |
|
|
| return img_copy |
|
|
| def predict_image(image, confidence, iou, imgsz): |
| if image is None: |
| return None, "#### Please upload an image to begin detection" |
|
|
| try: |
| img_bytes = io.BytesIO() |
| image.save(img_bytes, format='JPEG') |
| img_bytes.seek(0) |
|
|
| params = { |
| "conf": confidence, |
| "iou": iou, |
| "imgsz": imgsz |
| } |
|
|
| headers = {"Authorization": f"Bearer {API_KEY}"} |
| files = {"file": ("image.jpg", img_bytes, "image/jpeg")} |
|
|
| response = requests.post(API_URL, headers=headers, data=params, files=files, timeout=30) |
| response.raise_for_status() |
|
|
| result = response.json() |
| formatted_result = format_results(result) |
|
|
| detections = [] |
| if "images" in result and len(result["images"]) > 0: |
| detections = result["images"][0].get("results", []) |
|
|
| image_with_boxes = draw_boxes_on_image(image, detections) |
|
|
| return image_with_boxes, formatted_result |
|
|
| except requests.exceptions.Timeout: |
| return None, "#### Error: Request timeout. Please try again." |
| except requests.exceptions.ConnectionError: |
| return None, "#### Error: Unable to connect to detection service. Please check API configuration." |
| except requests.exceptions.HTTPError as e: |
| return None, f"#### Error: API returned status {e.response.status_code}" |
| except Exception as e: |
| return None, f"#### Error: {str(e)}" |
|
|
| def format_results(result): |
| if isinstance(result, dict): |
| output = "## Detection Results\n\n" |
|
|
| if "images" in result and len(result["images"]) > 0: |
| img_data = result["images"][0] |
| shape = img_data.get("shape", []) |
| detections = img_data.get("results", []) |
|
|
| output += f"**Image Size:** {shape[0]} x {shape[1]} (W x H)\n" |
| output += f"**Detections Found:** {len(detections)}\n\n" |
|
|
| speed = img_data.get("speed", {}) |
| if speed: |
| output += "\n### Performance Metrics\n" |
| output += "| Metric | Time (ms) |\n" |
| output += "|--------|----------|\n" |
| output += f"| Preprocess | {speed.get('preprocess', 'N/A')} |\n" |
| output += f"| Inference | {speed.get('inference', 'N/A')} |\n" |
| output += f"| Postprocess | {speed.get('postprocess', 'N/A')} |\n" |
|
|
| if detections: |
| output += "### Detected Objects\n" |
| output += "| Label | Class | Confidence |\n" |
| output += "|-------|-------|------------|\n" |
|
|
| for det in detections: |
| name = det.get("name", "Unknown") |
| class_id = det.get("class", "N/A") |
| conf = det.get("confidence", 0) |
| output += f"| {name} | {class_id} | {conf:.2%} |\n" |
|
|
| return output |
|
|
| return str(result) |
|
|
| dark_theme = gr.themes.Monochrome( |
| primary_hue="slate", |
| secondary_hue="slate", |
| ).set( |
| body_text_color="#e0e0e0", |
| background_fill_primary="#0f0f0f", |
| background_fill_secondary="#1a1a1a", |
| ) |
|
|
| with gr.Blocks( |
| title="YOLO Object Detection", |
| theme=dark_theme, |
| css=""" |
| footer {display: none !important;} |
| .gradio-container {border-radius: 12px;} |
| .gr-card {border-radius: 12px;} |
| .block {border-radius: 12px;} |
| .form {border-radius: 12px;} |
| button {border-radius: 12px;} |
| .gr-button {border-radius: 12px;} |
| #imageModal { |
| display: none; |
| position: fixed; |
| z-index: 10000; |
| left: 0; |
| top: 0; |
| width: 100%; |
| height: 100%; |
| background-color: rgba(0, 0, 0, 0.9); |
| animation: fadeIn 0.3s; |
| } |
| @keyframes fadeIn { |
| from {opacity: 0;} |
| to {opacity: 1;} |
| } |
| #modalImage { |
| position: absolute; |
| top: 50%; |
| left: 50%; |
| transform: translate(-50%, -50%); |
| max-width: 95%; |
| max-height: 95%; |
| object-fit: contain; |
| touch-action: pinch-zoom; |
| cursor: zoom-out; |
| } |
| .modal-open { |
| overflow: hidden; |
| } |
| .closeBtn { |
| position: absolute; |
| top: 20px; |
| right: 30px; |
| font-size: 40px; |
| font-weight: bold; |
| color: white; |
| cursor: pointer; |
| z-index: 10001; |
| } |
| .closeBtn:hover { |
| color: #bbb; |
| } |
| """ |
| ) as demo: |
| with gr.Column(): |
| gr.Markdown(""" |
| # YOLO Object Detection |
| Deploy your YOLO model with precision. Upload an image and adjust parameters to detect objects with custom inference settings. |
| """) |
|
|
| with gr.Row(): |
| with gr.Column(scale=1, min_width=400): |
| gr.Markdown("### Input") |
| image_input = gr.Image( |
| label="Image", |
| type="pil", |
| sources=["upload"], |
| interactive=True |
| ) |
| |
| test_images = get_test_images() |
| if test_images: |
| test_image_radio = gr.Radio( |
| choices=[img[1] for img in test_images], |
| label="Select test image", |
| info="Click to load" |
| ) |
| test_image_radio.change( |
| fn=lambda name: load_test_image(next((img[0] for img in test_images if img[1] == name), None)), |
| inputs=[test_image_radio], |
| outputs=[image_input] |
| ) |
| else: |
| gr.Markdown("No test images found. Add images to the 'images' folder.") |
|
|
| gr.Markdown("### Configuration") |
|
|
| confidence_slider = gr.Slider( |
| label="Confidence Threshold", |
| minimum=0.0, |
| maximum=1.0, |
| value=0.25, |
| step=0.01, |
| info="Detection confidence level" |
| ) |
|
|
| iou_slider = gr.Slider( |
| label="IOU Threshold", |
| minimum=0.0, |
| maximum=1.0, |
| value=0.7, |
| step=0.01, |
| info="Intersection over union threshold" |
| ) |
|
|
| imgsz_slider = gr.Slider( |
| label="Image Size", |
| minimum=320, |
| maximum=1280, |
| value=640, |
| step=32, |
| info="Inference image resolution" |
| ) |
|
|
| predict_btn = gr.Button( |
| "Detect Objects", |
| variant="primary", |
| size="lg", |
| scale=1 |
| ) |
|
|
| with gr.Column(scale=1, min_width=400): |
| gr.Markdown("### Results") |
|
|
| image_output = gr.Image( |
| label="Detections (Click to fullscreen)", |
| type="pil", |
| interactive=False, |
| scale=1 |
| ) |
|
|
| results_output = gr.Markdown( |
| value="Detection results will appear here.", |
| label="Detection Results" |
| ) |
|
|
| gr.HTML(""" |
| <div id="imageModal"> |
| <span class="closeBtn">×</span> |
| <img id="modalImage" src="" alt="Fullscreen Detection"> |
| </div> |
| <script> |
| const modal = document.getElementById('imageModal'); |
| const modalImg = document.getElementById('modalImage'); |
| const closeBtn = document.querySelector('.closeBtn'); |
| let touchStartX = 0; |
| let touchStartY = 0; |
| let scale = 1; |
| const observeImageChanges = () => { |
| const imageContainer = document.querySelector('[data-testid="image"]') || |
| document.querySelector('img[alt="Image"]'); |
| if (imageContainer) { |
| const images = imageContainer.querySelectorAll('img'); |
| images.forEach(img => { |
| if (img.src && !img.hasClickListener) { |
| img.style.cursor = 'pointer'; |
| img.addEventListener('click', (e) => { |
| if (e.target.src && !e.target.src.includes('data:image/svg')) { |
| modalImg.src = e.target.src; |
| modal.style.display = 'block'; |
| document.body.classList.add('modal-open'); |
| scale = 1; |
| modalImg.style.transform = 'translate(-50%, -50%) scale(1)'; |
| } |
| }); |
| img.hasClickListener = true; |
| } |
| }); |
| } |
| }; |
| setInterval(observeImageChanges, 500); |
| observeImageChanges(); |
| modal.addEventListener('click', (e) => { |
| if (e.target === modal) { |
| modal.style.display = 'none'; |
| document.body.classList.remove('modal-open'); |
| scale = 1; |
| } |
| }); |
| closeBtn.addEventListener('click', () => { |
| modal.style.display = 'none'; |
| document.body.classList.remove('modal-open'); |
| scale = 1; |
| }); |
| document.addEventListener('keydown', (e) => { |
| if (e.key === 'Escape' && modal.style.display === 'block') { |
| modal.style.display = 'none'; |
| document.body.classList.remove('modal-open'); |
| scale = 1; |
| } |
| }); |
| let lastDistance = 0; |
| modalImg.addEventListener('touchstart', (e) => { |
| if (e.touches.length === 2) { |
| const dx = e.touches[0].clientX - e.touches[1].clientX; |
| const dy = e.touches[0].clientY - e.touches[1].clientY; |
| lastDistance = Math.sqrt(dx * dx + dy * dy); |
| } |
| touchStartX = e.touches[0].clientX; |
| touchStartY = e.touches[0].clientY; |
| }); |
| modalImg.addEventListener('touchmove', (e) => { |
| if (e.touches.length === 2) { |
| const dx = e.touches[0].clientX - e.touches[1].clientX; |
| const dy = e.touches[0].clientY - e.touches[1].clientY; |
| const distance = Math.sqrt(dx * dx + dy * dy); |
| const scaleChange = distance / lastDistance; |
| scale = Math.max(1, Math.min(scale * scaleChange, 4)); |
| modalImg.style.transform = `translate(-50%, -50%) scale(${scale})`; |
| lastDistance = distance; |
| } |
| }); |
| modalImg.addEventListener('touchend', () => { |
| lastDistance = 0; |
| }); |
| </script> |
| """) |
|
|
| predict_btn.click( |
| fn=predict_image, |
| inputs=[image_input, confidence_slider, iou_slider, imgsz_slider], |
| outputs=[image_output, results_output] |
| ) |
|
|
| image_input.change( |
| fn=predict_image, |
| inputs=[image_input, confidence_slider, iou_slider, imgsz_slider], |
| outputs=[image_output, results_output] |
| ) |
|
|
| confidence_slider.change( |
| fn=predict_image, |
| inputs=[image_input, confidence_slider, iou_slider, imgsz_slider], |
| outputs=[image_output, results_output] |
| ) |
|
|
| iou_slider.change( |
| fn=predict_image, |
| inputs=[image_input, confidence_slider, iou_slider, imgsz_slider], |
| outputs=[image_output, results_output] |
| ) |
|
|
| imgsz_slider.change( |
| fn=predict_image, |
| inputs=[image_input, confidence_slider, iou_slider, imgsz_slider], |
| outputs=[image_output, results_output] |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch(share=False, show_error=True) |
