Spaces:
Running
Running
| import gradio as gr | |
| import requests | |
| import json | |
| import base64 | |
| from PIL import Image, ImageDraw, ImageFont | |
| import io | |
| def process_with_openrouter(image, prompt, api_key, model="qwen/qwen2.5-vl-32b-instruct", temperature=0.5): | |
| """Process image with OpenRouter API for object detection""" | |
| if not api_key: | |
| return "Please enter your OpenRouter API key", "error" | |
| if image is None: | |
| return "Please upload an image", "error" | |
| try: | |
| buffered = io.BytesIO() | |
| image.save(buffered, format="PNG") | |
| img_base64 = base64.b64encode(buffered.getvalue()).decode() | |
| headers = { | |
| "Authorization": f"Bearer {api_key}", | |
| "Content-Type": "application/json" | |
| } | |
| data = { | |
| "model": model, | |
| "messages": [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": prompt}, | |
| { | |
| "type": "image_url", | |
| "image_url": {"url": f"data:image/png;base64,{img_base64}"} | |
| } | |
| ] | |
| } | |
| ], | |
| "temperature": temperature | |
| } | |
| response = requests.post( | |
| "https://openrouter.ai/api/v1/chat/completions", | |
| headers=headers, | |
| json=data, | |
| timeout=60 | |
| ) | |
| if response.status_code == 200: | |
| result = response.json() | |
| content = result['choices'][0]['message']['content'] | |
| if '```json' in content: | |
| content = content.split('```json')[1].split('```')[0].strip() | |
| elif '```' in content: | |
| content = content.split('```')[1].split('```')[0].strip() | |
| return content, None | |
| else: | |
| return f"Error: {response.status_code} - {response.text}", "error" | |
| except Exception as e: | |
| return f"Error processing request: {str(e)}", "error" | |
| def draw_bounding_boxes(image, detections): | |
| """Draw bounding boxes with class names only, same color per class""" | |
| if not detections or len(detections) == 0: | |
| return image | |
| annotated_image = image.copy() | |
| draw = ImageDraw.Draw(annotated_image) | |
| try: | |
| font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 16) | |
| except: | |
| font = ImageFont.load_default() | |
| # Predefined colors for different classes | |
| class_colors = { | |
| "Class I": "#FF0000", # Red | |
| "Class II": "#00FF00", # Green | |
| "Class III": "#0000FF", # Blue | |
| "Class IV": "#FFFF00", # Yellow | |
| "Class V": "#FF00FF", # Magenta | |
| "Class VI": "#00FFFF", # Cyan | |
| "Class VII": "#FFA500", # Orange | |
| "Class VIII": "#800080", # Purple | |
| "Class IX": "#008000", # Dark Green | |
| "Class X": "#FF1493", # Deep Pink | |
| } | |
| # Fallback colors if more than 10 classes | |
| fallback_colors = ["#8B4513", "#2F4F4F", "#DC143C", "#00CED1", "#FF4500", "#DA70D6", "#32CD32", "#FF6347"] | |
| for i, detection in enumerate(detections): | |
| if all(key in detection for key in ['x', 'y', 'width', 'height']): | |
| x = detection['x'] * image.width | |
| y = detection['y'] * image.height | |
| width = detection['width'] * image.width | |
| height = detection['height'] * image.height | |
| # Get class name - this is what we'll display | |
| class_name = detection.get('class', f'Class {i+1}') | |
| x1, y1 = int(x), int(y) | |
| x2, y2 = int(x + width), int(y + height) | |
| x1 = max(0, min(x1, image.width)) | |
| y1 = max(0, min(y1, image.height)) | |
| x2 = max(0, min(x2, image.width)) | |
| y2 = max(0, min(y2, image.height)) | |
| # Get consistent color for this class | |
| if class_name in class_colors: | |
| color = class_colors[class_name] | |
| else: | |
| # Use hash of class name to get consistent color | |
| color_index = hash(class_name) % len(fallback_colors) | |
| color = fallback_colors[color_index] | |
| # Draw bounding box | |
| draw.rectangle([x1, y1, x2, y2], outline=color, width=4) | |
| # Calculate label size | |
| text_bbox = draw.textbbox((0, 0), class_name, font=font) | |
| text_width = text_bbox[2] - text_bbox[0] | |
| text_height = text_bbox[3] - text_bbox[1] | |
| # Position label above the box, or below if no space above | |
| if y1 - text_height - 6 >= 0: | |
| label_y = y1 - text_height - 6 | |
| else: | |
| label_y = y2 + 4 | |
| label_x = x1 | |
| # Ensure label stays within image bounds | |
| if label_x + text_width + 4 > image.width: | |
| label_x = image.width - text_width - 4 | |
| # Draw label background | |
| draw.rectangle( | |
| [label_x - 2, label_y - 2, label_x + text_width + 2, label_y + text_height + 2], | |
| fill=color, | |
| outline=color | |
| ) | |
| # Draw class name | |
| draw.text((label_x, label_y), class_name, fill="white", font=font) | |
| return annotated_image | |
| def create_detection_prompt(class_descriptions, confidence_threshold=0.5, detection_mode="specific"): | |
| """Create a detection prompt for class descriptions with condition checking""" | |
| if isinstance(class_descriptions, str): | |
| class_descriptions = [cls.strip() for cls in class_descriptions.split('\n') if cls.strip()] | |
| # Build detection instructions | |
| if detection_mode == "specific": | |
| condition_text = "ONLY detect objects that match these class descriptions and their conditions. Ignore all other objects:" | |
| elif detection_mode == "include": | |
| condition_text = "Detect objects matching these class descriptions AND any other objects you can identify:" | |
| else: # "exclude" | |
| condition_text = "Detect all objects EXCEPT those matching these class descriptions. Avoid detecting:" | |
| # Format each class description | |
| class_specs = [] | |
| for i, description in enumerate(class_descriptions, 1): | |
| # Parse class name and description if formatted as "Class Name: description" | |
| if ':' in description: | |
| class_name, class_desc = description.split(':', 1) | |
| class_name = class_name.strip() | |
| class_desc = class_desc.strip() | |
| class_specs.append(f"Class {i} ({class_name}): {class_desc}") | |
| else: | |
| class_specs.append(f"Class {i}: {description}") | |
| classes_text = "\n".join(class_specs) if class_specs else "No class descriptions provided" | |
| prompt = f"""{condition_text} | |
| {classes_text} | |
| Detection Instructions: | |
| - Analyze each object against the class descriptions above | |
| - Check if objects meet the specified conditions for each class | |
| - Only include detections with confidence above {confidence_threshold} | |
| - Assign objects to the most appropriate class based on the descriptions | |
| SCALE/RULER DETECTION FOR CRACK MEASUREMENT: | |
| - First look for scales, rulers, measurement tools, or reference objects in the image | |
| - If found, identify the scale markings and determine the measurement reference | |
| - Use the scale to calculate actual crack widths in millimeters or appropriate units | |
| - For crack classifications, measure crack width using the identified scale | |
| - Include actual measurements in your analysis (e.g., "2.5mm crack width based on ruler scale") | |
| - If no scale is visible, estimate crack width relative to common objects or provide qualitative assessment | |
| Output a JSON list where each entry contains: | |
| - "x": normalized x coordinate (0-1) of top-left corner | |
| - "y": normalized y coordinate (0-1) of top-left corner | |
| - "width": normalized width (0-1) of the bounding box | |
| - "height": normalized height (0-1) of the bounding box | |
| - "label": brief description with confidence score | |
| - "confidence": confidence score (0-1) | |
| - "class": the assigned class name (e.g., "Class I", "Class II", etc.) | |
| - "description": why this object matches the class criteria | |
| - "class_number": the class number from the list above (1, 2, 3, etc.) | |
| - "measured_width": actual crack width measurement if scale is available (e.g., "2.5mm", "1.2cm") | |
| - "measurement_method": how the measurement was obtained (e.g., "ruler scale", "coin reference", "estimated") | |
| Example format: | |
| [{{"x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4, "label": "Structural crack (0.92)", "confidence": 0.92, "class": "Class I", "description": "Crack width exceeds 2mm threshold based on ruler measurement", "class_number": 1, "measured_width": "2.5mm", "measurement_method": "ruler scale"}}]""" | |
| return prompt | |
| def create_interface(): | |
| """Create the Gradio interface for object detection""" | |
| with gr.Blocks(title="Class-Based Object Detection", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# 🔍 Class-Based Object Detection with Descriptions") | |
| gr.Markdown("Define classes with descriptions and conditions. Objects will be classified and annotated with class names only.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("## ⚙️ Configuration") | |
| api_key = gr.Textbox( | |
| label="OpenRouter API Key", | |
| placeholder="Enter your OpenRouter API key...", | |
| type="password" | |
| ) | |
| with gr.Row(): | |
| use_preset = gr.Radio( | |
| choices=["Preset Model", "Custom Model"], | |
| value="Preset Model", | |
| label="Model Selection", | |
| info="Choose preset or enter custom OpenRouter model" | |
| ) | |
| model_preset = gr.Dropdown( | |
| choices=[ | |
| "qwen/qwen2.5-vl-32b-instruct", | |
| "qwen/qwen-vl-max", | |
| "openai/gpt-5-chat", | |
| "openai/gpt-5-mini", | |
| "anthropic/claude-opus-4.1", | |
| "x-ai/grok-4", | |
| "google/gemini-2.5-pro", | |
| "google/gemini-1.5-pro", | |
| "google/gemini-1.5-flash", | |
| "anthropic/claude-3.5-sonnet", | |
| "openai/gpt-4o", | |
| "openai/gpt-4o-mini" | |
| ], | |
| value="qwen/qwen2.5-vl-32b-instruct", | |
| label="Preset Models", | |
| info="Select from popular OpenRouter models", | |
| visible=True | |
| ) | |
| custom_model_input = gr.Textbox( | |
| label="Custom Model ID", | |
| placeholder="Enter any OpenRouter model ID (e.g., google/gemini-1.5-flash, anthropic/claude-3-haiku)", | |
| visible=False, | |
| info="Copy model IDs from openrouter.ai/models" | |
| ) | |
| detection_mode = gr.Radio( | |
| choices=[ | |
| ("Detect Only These Classes", "specific"), | |
| ("Include These Classes + Others", "include"), | |
| ("Exclude These Classes", "exclude") | |
| ], | |
| value="specific", | |
| label="Detection Mode", | |
| info="How to handle the specified class descriptions" | |
| ) | |
| class_descriptions = gr.Textbox( | |
| label="Class Descriptions", | |
| placeholder="""Define each class with its description and conditions, e.g.: | |
| Severe Cracks: Crack width more than 2mm (use ruler/scale if present for measurement) | |
| Minor Cracks: Crack width 0.5-2mm (measure using visible scale) | |
| Rust Damage: Rust spots larger than 5cm in diameter | |
| Concrete Spalling: Concrete spalling deeper than 1cm | |
| Paint Defects: Paint peeling areas greater than 10cm²""", | |
| value="""Severe Cracks: Crack width more than 2mm (use ruler/scale if present for measurement) | |
| Minor Cracks: Crack width 0.5-2mm (measure using visible scale) | |
| Rust Damage: Rust spots larger than 5cm in diameter""", | |
| lines=8, | |
| info="Enter class descriptions, one per line. Format: 'Class Name: Description' or just 'Description'" | |
| ) | |
| confidence_threshold = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.5, | |
| step=0.05, | |
| label="Confidence Threshold", | |
| info="Minimum confidence for detection" | |
| ) | |
| temperature = gr.Slider( | |
| minimum=0, | |
| maximum=1, | |
| value=0.3, | |
| step=0.05, | |
| label="Temperature", | |
| info="Lower values for more consistent results" | |
| ) | |
| image_input = gr.Image( | |
| type="pil", | |
| label="Upload Image for Detection" | |
| ) | |
| detect_btn = gr.Button("🚀 Detect Objects", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| gr.Markdown("## 📊 Detection Results") | |
| annotated_image = gr.Image( | |
| label="Detected Objects", | |
| type="pil" | |
| ) | |
| detection_results = gr.Textbox( | |
| label="Detection Details (JSON)", | |
| lines=10, | |
| show_copy_button=True | |
| ) | |
| detection_summary = gr.Textbox( | |
| label="Detection Summary", | |
| lines=3 | |
| ) | |
| # Show/hide model input based on selection | |
| def update_model_visibility(use_preset_val): | |
| if use_preset_val == "Custom Model": | |
| return gr.update(visible=False), gr.update(visible=True) | |
| else: | |
| return gr.update(visible=True), gr.update(visible=False) | |
| use_preset.change( | |
| update_model_visibility, | |
| inputs=[use_preset], | |
| outputs=[model_preset, custom_model_input] | |
| ) | |
| def process_detection(image, class_desc, conf_threshold, api_key_val, use_preset_val, model_preset_val, custom_model_val, temp_val, mode_val): | |
| if not api_key_val: | |
| return None, "❌ Please enter your OpenRouter API key", "No API key provided" | |
| if image is None: | |
| return None, "❌ Please upload an image", "No image uploaded" | |
| if not class_desc or not class_desc.strip(): | |
| return None, "❌ Please enter at least one class description", "No class descriptions provided" | |
| # Determine which model to use | |
| if use_preset_val == "Custom Model": | |
| if not custom_model_val or custom_model_val.strip() == "": | |
| return None, "❌ Please enter a custom model ID", "Custom model required" | |
| final_model = custom_model_val.strip() | |
| else: | |
| final_model = model_preset_val | |
| try: | |
| prompt = create_detection_prompt(class_desc, conf_threshold, mode_val) | |
| result, error = process_with_openrouter(image, prompt, api_key_val, final_model, temp_val) | |
| if error: | |
| return None, f"❌ Error: {result}", "Detection failed" | |
| detections = json.loads(result) | |
| if isinstance(detections, list) and len(detections) > 0: | |
| annotated_img = draw_bounding_boxes(image, detections) | |
| filtered_detections = [d for d in detections if d.get('confidence', 1.0) >= conf_threshold] | |
| mode_descriptions = { | |
| "specific": "Detecting only objects matching class descriptions", | |
| "include": "Including specified classes + other objects", | |
| "exclude": "Excluding objects matching class descriptions" | |
| } | |
| summary_text = f"✅ {mode_descriptions.get(mode_val, 'Detection')} - Found {len(filtered_detections)} objects\n🤖 Model: {final_model}" | |
| if filtered_detections: | |
| # Group by class and show counts | |
| class_counts = {} | |
| for det in filtered_detections: | |
| class_name = det.get('class', 'unknown') | |
| description = det.get('description', '') | |
| confidence = det.get('confidence', 1.0) | |
| if class_name not in class_counts: | |
| class_counts[class_name] = { | |
| 'count': 0, | |
| 'avg_confidence': 0, | |
| 'descriptions': [] | |
| } | |
| class_counts[class_name]['count'] += 1 | |
| class_counts[class_name]['avg_confidence'] += confidence | |
| if description and description not in class_counts[class_name]['descriptions']: | |
| class_counts[class_name]['descriptions'].append(description) | |
| summary_text += "\n\nClass Detection Results:" | |
| for class_name, data in class_counts.items(): | |
| avg_conf = data['avg_confidence'] / data['count'] | |
| summary_text += f"\n• {class_name}: {data['count']} detected (avg conf: {avg_conf:.2f})" | |
| return annotated_img, json.dumps(filtered_detections, indent=2), summary_text | |
| else: | |
| return image, "No objects detected matching class descriptions", "No detections matching criteria above confidence threshold" | |
| except json.JSONDecodeError: | |
| return None, f"❌ Invalid JSON response: {result}", "JSON parsing failed" | |
| except Exception as e: | |
| return None, f"❌ Error: {str(e)}", "Processing error" | |
| detect_btn.click( | |
| process_detection, | |
| inputs=[image_input, class_descriptions, confidence_threshold, api_key, use_preset, model_preset, custom_model_input, temperature, detection_mode], | |
| outputs=[annotated_image, detection_results, detection_summary] | |
| ) | |
| gr.Markdown(""" | |
| ## 💡 Usage Tips | |
| - **Specific Mode**: Only detect objects matching your class descriptions | |
| - **Include Mode**: Detect your specified classes plus any other objects found | |
| - **Exclude Mode**: Detect everything except objects matching your class descriptions | |
| ### 🏷️ Class Definition | |
| **Format Options:** | |
| 1. `Class Name: Description` - e.g., "Severe Cracks: Crack width more than 2mm" | |
| 2. `Description only` - Will be automatically assigned as "Class I", "Class II", etc. | |
| **Annotation Behavior:** | |
| - Images show only class names (e.g., "Class I", "Class II") | |
| - Same class = same color throughout the image | |
| - Clean, simple visual identification | |
| ### 🤖 Model Selection | |
| **Default Models (Recommended):** | |
| - `qwen/qwen2.5-vl-32b-instruct` - Advanced Qwen vision model optimized for detailed analysis (Default) | |
| - `qwen/qwen-vl-max` - Premium Qwen vision model with maximum capabilities | |
| - `openai/gpt-5-chat` - Latest GPT-5 with advanced vision capabilities | |
| - `openai/gpt-5-mini` - Faster, efficient GPT-5 variant | |
| - `anthropic/claude-opus-4.1` - Next-gen Claude with superior reasoning | |
| - `x-ai/grok-4` - Advanced Grok model with detailed analysis | |
| **Custom Models**: Enter any OpenRouter model ID from [openrouter.ai/models](https://openrouter.ai/models) | |
| ### Example Class Descriptions: | |
| ``` | |
| Severe Cracks: Crack width more than 2mm (use ruler/scale for measurement) | |
| Minor Cracks: Crack width 0.5-2mm (measure using visible scale) | |
| Rust Damage: Rust spots larger than 5cm in diameter | |
| Concrete Spalling: Concrete spalling deeper than 1cm | |
| Paint Defects: Paint peeling areas greater than 10cm² | |
| Water Damage: Water damage stains larger than 15cm | |
| ``` | |
| ### 📏 Scale-Based Measurement: | |
| - **Automatic Scale Detection**: The system looks for rulers, measuring tools, or reference objects | |
| - **Precise Measurements**: When scales are found, actual crack widths are calculated | |
| - **Measurement Methods**: Supports rulers, crack gauges, coins, or other reference objects | |
| - **Enhanced Classification**: More accurate class assignment based on measured dimensions | |
| - Enter one class description per line | |
| - Be specific about conditions and measurements | |
| - Objects will be classified and labeled with class names only | |
| - Adjust confidence threshold to filter weak detections | |
| - Get your API key from [openrouter.ai](https://openrouter.ai/) | |
| """) | |
| return demo | |
| if __name__ == "__main__": | |
| print("🚀 Starting Object Detection App...") | |
| demo = create_interface() | |
| demo.launch(share=False, inbrowser=True) |