| # import gradio as gr | |
| # import numpy as np | |
| # import torch | |
| # from PIL import Image | |
| # import matplotlib.pyplot as plt | |
| # from transformers import pipeline | |
| # import warnings | |
| # from io import BytesIO | |
| # import importlib.util | |
| # import os | |
| # import openai | |
| # # Suppress warnings | |
| # warnings.filterwarnings("ignore") | |
| # # Set up OpenAI API key | |
| # api_key = os.getenv('OPENAI_API_KEY') | |
| # if not api_key: | |
| # print("No OpenAI API key found - will use simple keyword extraction") | |
| # elif not api_key.startswith("sk-proj-") and not api_key.startswith("sk-"): | |
| # print("API key found but doesn't look correct") | |
| # elif api_key.strip() != api_key: | |
| # print("API key has leading or trailing whitespace - please fix it.") | |
| # else: | |
| # print("OpenAI API key found and looks good!") | |
| # openai.api_key = api_key | |
| # # Global variables for models | |
| # detector = None | |
| # sam_predictor = None | |
| # def load_detector(): | |
| # """Load the OWL-ViT detector once and cache it.""" | |
| # global detector | |
| # if detector is None: | |
| # print("Loading OWL-ViT model...") | |
| # detector = pipeline( | |
| # model="google/owlv2-base-patch16-ensemble", | |
| # task="zero-shot-object-detection", | |
| # device=0 if torch.cuda.is_available() else -1 | |
| # ) | |
| # print("OWL-ViT model loaded successfully!") | |
| # def install_sam2_if_needed(): | |
| # """Check if SAM2 is installed, and install it if needed.""" | |
| # if importlib.util.find_spec("sam2") is not None: | |
| # print("SAM2 is already installed.") | |
| # return True | |
| # try: | |
| # import subprocess | |
| # import sys | |
| # print("Installing SAM2 from GitHub...") | |
| # subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/facebookresearch/sam2.git"]) | |
| # print("SAM2 installed successfully.") | |
| # return True | |
| # except Exception as e: | |
| # print(f"Error installing SAM2: {e}") | |
| # return False | |
| # def load_sam_predictor(): | |
| # """Load SAM2 predictor if available.""" | |
| # global sam_predictor | |
| # if sam_predictor is None: | |
| # if install_sam2_if_needed(): | |
| # try: | |
| # from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| # print("Loading SAM2 model...") | |
| # sam_predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large") | |
| # device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # sam_predictor.model.to(device) | |
| # print(f"SAM2 model loaded successfully on {device}!") | |
| # return True | |
| # except Exception as e: | |
| # print(f"Error loading SAM2: {e}") | |
| # return False | |
| # return sam_predictor is not None | |
| # def detect_objects_owlv2(text_query, image, threshold=0.1): | |
| # """Detect objects using OWL-ViT.""" | |
| # try: | |
| # load_detector() | |
| # if isinstance(image, np.ndarray): | |
| # image = Image.fromarray(image) | |
| # # Clean up the text query | |
| # query_terms = [term.strip() for term in text_query.split(',') if term.strip()] | |
| # if not query_terms: | |
| # query_terms = ["object"] | |
| # print(f"Detecting: {query_terms}") | |
| # predictions = detector(image, candidate_labels=query_terms) | |
| # detections = [] | |
| # for pred in predictions: | |
| # if pred['score'] >= threshold: | |
| # bbox = pred['box'] | |
| # width, height = image.size | |
| # normalized_bbox = [ | |
| # bbox['xmin'] / width, | |
| # bbox['ymin'] / height, | |
| # bbox['xmax'] / width, | |
| # bbox['ymax'] / height | |
| # ] | |
| # detection = { | |
| # 'label': pred['label'], | |
| # 'bbox': normalized_bbox, | |
| # 'score': pred['score'] | |
| # } | |
| # detections.append(detection) | |
| # return detections, image | |
| # except Exception as e: | |
| # print(f"Detection error: {e}") | |
| # return [], image | |
| # def generate_masks_sam2(detections, image): | |
| # """Generate segmentation masks using SAM2.""" | |
| # try: | |
| # if not load_sam_predictor(): | |
| # print("SAM2 not available, skipping mask generation") | |
| # return detections | |
| # if isinstance(image, np.ndarray): | |
| # image = Image.fromarray(image) | |
| # image_np = np.array(image.convert("RGB")) | |
| # H, W = image_np.shape[:2] | |
| # # Set image for SAM2 | |
| # sam_predictor.set_image(image_np) | |
| # # Convert normalized bboxes to pixel coordinates | |
| # input_boxes = [] | |
| # for det in detections: | |
| # x1, y1, x2, y2 = det['bbox'] | |
| # input_boxes.append([int(x1 * W), int(y1 * H), int(x2 * W), int(y2 * H)]) | |
| # if not input_boxes: | |
| # return detections | |
| # input_boxes = np.array(input_boxes) | |
| # print(f"Generating masks for {len(input_boxes)} detections...") | |
| # with torch.inference_mode(): | |
| # device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # if device == "cuda": | |
| # with torch.autocast("cuda", dtype=torch.bfloat16): | |
| # masks, scores, _ = sam_predictor.predict( | |
| # point_coords=None, | |
| # point_labels=None, | |
| # box=input_boxes, | |
| # multimask_output=False | |
| # ) | |
| # else: | |
| # masks, scores, _ = sam_predictor.predict( | |
| # point_coords=None, | |
| # point_labels=None, | |
| # box=input_boxes, | |
| # multimask_output=False | |
| # ) | |
| # # Add masks to detections | |
| # results = [] | |
| # for i, det in enumerate(detections): | |
| # new_det = det.copy() | |
| # mask = masks[i] | |
| # if mask.ndim == 3: | |
| # mask = mask[0] # Remove batch dimension if present | |
| # new_det['mask'] = mask.astype(np.uint8) | |
| # results.append(new_det) | |
| # print(f"Successfully generated {len(results)} masks") | |
| # return results | |
| # except Exception as e: | |
| # print(f"SAM2 mask generation error: {e}") | |
| # return detections | |
| # def visualize_detections_with_masks(image, detections_with_masks, show_labels=True, show_boxes=True): | |
| # """ | |
| # Visualize the detections with their segmentation masks. | |
| # Returns PIL Image instead of showing plot. | |
| # """ | |
| # # Load the image | |
| # if isinstance(image, np.ndarray): | |
| # image = Image.fromarray(image) | |
| # image_np = np.array(image.convert("RGB")) | |
| # # Get image dimensions | |
| # height, width = image_np.shape[:2] | |
| # # Create figure | |
| # fig = plt.figure(figsize=(12, 8)) | |
| # plt.imshow(image_np) | |
| # # Define colors for different instances | |
| # colors = plt.cm.tab10(np.linspace(0, 1, 10)) | |
| # # Plot each detection | |
| # for i, detection in enumerate(detections_with_masks): | |
| # # Get bbox, mask, label, and score | |
| # bbox = detection['bbox'] | |
| # label = detection['label'] | |
| # score = detection['score'] | |
| # # Convert normalized bbox to pixel coordinates | |
| # x1, y1, x2, y2 = bbox | |
| # x1_px, y1_px = int(x1 * width), int(y1 * height) | |
| # x2_px, y2_px = int(x2 * width), int(y2 * height) | |
| # # Color for this instance | |
| # color = colors[i % len(colors)] | |
| # # Display mask with transparency if available | |
| # if 'mask' in detection: | |
| # mask = detection['mask'] | |
| # mask_color = np.zeros((height, width, 4), dtype=np.float32) | |
| # mask_color[mask > 0] = [color[0], color[1], color[2], 0.5] | |
| # plt.imshow(mask_color) | |
| # # Draw bounding box if requested | |
| # if show_boxes: | |
| # rect = plt.Rectangle((x1_px, y1_px), x2_px - x1_px, y2_px - y1_px, | |
| # fill=False, edgecolor=color, linewidth=2) | |
| # plt.gca().add_patch(rect) | |
| # # Add label and score if requested | |
| # if show_labels: | |
| # plt.text(x1_px, y1_px - 5, f"{label}: {score:.2f}", | |
| # color='white', bbox=dict(facecolor=color, alpha=0.8), fontsize=10) | |
| # plt.axis('off') | |
| # plt.tight_layout() | |
| # # Convert to PIL Image using the correct method | |
| # buf = BytesIO() | |
| # plt.savefig(buf, format='png', bbox_inches='tight', dpi=150) | |
| # plt.close(fig) | |
| # buf.seek(0) | |
| # result_image = Image.open(buf) | |
| # return result_image | |
| # def visualize_detections(image, detections, show_labels=True): | |
| # """ | |
| # Visualize object detections with bounding boxes only. | |
| # Returns PIL Image instead of showing plot. | |
| # """ | |
| # # Load the image | |
| # if isinstance(image, np.ndarray): | |
| # image = Image.fromarray(image) | |
| # image_np = np.array(image.convert("RGB")) | |
| # # Get image dimensions | |
| # height, width = image_np.shape[:2] | |
| # # Create figure | |
| # fig = plt.figure(figsize=(12, 8)) | |
| # plt.imshow(image_np) | |
| # # If we have detections, draw them | |
| # if detections: | |
| # # Define colors for different instances | |
| # colors = plt.cm.tab10(np.linspace(0, 1, 10)) | |
| # # Plot each detection | |
| # for i, detection in enumerate(detections): | |
| # # Get bbox, label, and score | |
| # bbox = detection['bbox'] | |
| # label = detection['label'] | |
| # score = detection['score'] | |
| # # Convert normalized bbox to pixel coordinates | |
| # x1, y1, x2, y2 = bbox | |
| # x1_px, y1_px = int(x1 * width), int(y1 * height) | |
| # x2_px, y2_px = int(x2 * width), int(y2 * height) | |
| # # Color for this instance | |
| # color = colors[i % len(colors)] | |
| # # Draw bounding box | |
| # rect = plt.Rectangle((x1_px, y1_px), x2_px - x1_px, y2_px - y1_px, | |
| # fill=False, edgecolor=color, linewidth=2) | |
| # plt.gca().add_patch(rect) | |
| # # Add label and score if requested | |
| # if show_labels: | |
| # plt.text(x1_px, y1_px - 5, f"{label}: {score:.2f}", | |
| # color='white', bbox=dict(facecolor=color, alpha=0.8), fontsize=10) | |
| # # Set title | |
| # plt.title(f'Object Detection Results ({len(detections)} objects found)', fontsize=14, pad=20) | |
| # plt.axis('off') | |
| # plt.tight_layout() | |
| # # Convert to PIL Image | |
| # buf = BytesIO() | |
| # plt.savefig(buf, format='png', bbox_inches='tight', dpi=150) | |
| # plt.close(fig) | |
| # buf.seek(0) | |
| # result_image = Image.open(buf) | |
| # return result_image | |
| # def get_optimized_prompt(query_text): | |
| # """ | |
| # Use OpenAI to convert natural language query into optimal detection prompt. | |
| # Falls back to simple extraction if OpenAI is not available. | |
| # """ | |
| # if not query_text.strip(): | |
| # return "object" | |
| # # Try OpenAI first if API key is available | |
| # if hasattr(openai, 'api_key') and openai.api_key: | |
| # try: | |
| # response = openai.chat.completions.create( | |
| # model="gpt-3.5-turbo", | |
| # messages=[{ | |
| # "role": "system", | |
| # "content": """You are an expert at converting natural language queries into precise object detection terms. | |
| # RULES: | |
| # 1. Return ONLY 1-2 words maximum that describe the object to detect | |
| # 2. Use the exact object name from the user's query | |
| # 3. For people: use "person" | |
| # 4. For vehicles: use "car", "truck", "bicycle" | |
| # 5. Do NOT include counting words, articles, or explanations | |
| # 6. Examples: | |
| # - "How many cacao fruits are there?" → "cacao fruit" | |
| # - "Count the corn in the field" → "corn" | |
| # - "Find all people" → "person" | |
| # - "How many cacao pods?" → "cacao pod" | |
| # - "Detect cars" → "car" | |
| # - "Count bananas" → "banana" | |
| # - "How many apples?" → "apple" | |
| # Return ONLY the object name, nothing else.""" | |
| # }, { | |
| # "role": "user", | |
| # "content": query_text | |
| # }], | |
| # temperature=0.0, # Make it deterministic | |
| # max_tokens=5 # Force brevity | |
| # ) | |
| # llm_result = response.choices[0].message.content.strip().lower() | |
| # # Extra safety: take only first 2 words | |
| # words = llm_result.split()[:2] | |
| # final_result = " ".join(words) | |
| # print(f"🤖 OpenAI suggested prompt: '{final_result}'") | |
| # return final_result | |
| # except Exception as e: | |
| # print(f"OpenAI error: {e}, falling back to keyword extraction") | |
| # # Fallback to simple keyword extraction (no hardcoded fruits) | |
| # print("🔤 Using keyword extraction (no OpenAI)") | |
| # query_lower = query_text.lower().replace("?", "").strip() | |
| # # Look for common patterns and extract object names | |
| # if "how many" in query_lower: | |
| # parts = query_lower.split("how many") | |
| # if len(parts) > 1: | |
| # remaining = parts[1].strip() | |
| # remaining = remaining.replace("are", "").replace("in", "").replace("the", "").replace("image", "").replace("there", "").strip() | |
| # # Take first meaningful word(s) | |
| # words = remaining.split()[:2] | |
| # search_terms = " ".join(words) if words else "object" | |
| # else: | |
| # search_terms = "object" | |
| # elif "count" in query_lower: | |
| # parts = query_lower.split("count") | |
| # if len(parts) > 1: | |
| # remaining = parts[1].strip() | |
| # remaining = remaining.replace("the", "").replace("in", "").replace("image", "").strip() | |
| # words = remaining.split()[:2] | |
| # search_terms = " ".join(words) if words else "object" | |
| # else: | |
| # search_terms = "object" | |
| # elif "find" in query_lower: | |
| # parts = query_lower.split("find") | |
| # if len(parts) > 1: | |
| # remaining = parts[1].strip() | |
| # remaining = remaining.replace("all", "").replace("the", "").replace("in", "").replace("image", "").strip() | |
| # words = remaining.split()[:2] | |
| # search_terms = " ".join(words) if words else "object" | |
| # else: | |
| # search_terms = "object" | |
| # else: | |
| # # Extract first 1-2 meaningful words from the query | |
| # words = query_lower.split() | |
| # meaningful_words = [w for w in words if w not in ["how", "many", "are", "in", "the", "image", "find", "count", "detect", "there", "this", "that", "a", "an"]] | |
| # search_terms = " ".join(meaningful_words[:2]) if meaningful_words else "object" | |
| # return search_terms | |
| # def is_count_query(text): | |
| # """Check if the query is asking for counting.""" | |
| # count_keywords = ["how many", "count", "number of", "total"] | |
| # return any(keyword in text.lower() for keyword in count_keywords) | |
| # def detection_pipeline(query_text, image, threshold, use_sam): | |
| # """Main detection pipeline.""" | |
| # if image is None: | |
| # return None, "⚠️ Please upload an image first!" | |
| # try: | |
| # # Use OpenAI or fallback to get optimized search terms | |
| # search_terms = get_optimized_prompt(query_text) | |
| # print(f"Processing query: '{query_text}' -> searching for: '{search_terms}'") | |
| # # Run object detection | |
| # detections, processed_image = detect_objects_owlv2(search_terms, image, threshold) | |
| # print(f"Found {len(detections)} detections") | |
| # for i, det in enumerate(detections): | |
| # print(f"Detection {i+1}: {det['label']} (score: {det['score']:.3f})") | |
| # # Generate masks if requested | |
| # if use_sam and detections: | |
| # print("Generating SAM2 masks...") | |
| # detections = generate_masks_sam2(detections, processed_image) | |
| # # Create visualization using your proven functions | |
| # print("Creating visualization...") | |
| # if use_sam and detections and 'mask' in detections[0]: | |
| # result_image = visualize_detections_with_masks( | |
| # processed_image, | |
| # detections, | |
| # show_labels=True, | |
| # show_boxes=True | |
| # ) | |
| # print("Created visualization with masks") | |
| # else: | |
| # result_image = visualize_detections( | |
| # processed_image, | |
| # detections, | |
| # show_labels=True | |
| # ) | |
| # print("Created visualization with bounding boxes only") | |
| # # Make sure we have a valid result image | |
| # if result_image is None: | |
| # print("Warning: result_image is None, returning original image") | |
| # result_image = processed_image | |
| # # Generate summary | |
| # count = len(detections) | |
| # summary_parts = [] | |
| # summary_parts.append(f"🗣️ **Original Query**: '{query_text}'") | |
| # summary_parts.append(f"🤖 **AI-Optimized Search**: '{search_terms}'") | |
| # summary_parts.append(f"⚙️ **Threshold**: {threshold}") | |
| # summary_parts.append(f"🎭 **SAM2 Segmentation**: {'Enabled' if use_sam else 'Disabled'}") | |
| # if count > 0: | |
| # if is_count_query(query_text): | |
| # summary_parts.append(f"🔢 **Answer: {count} {search_terms}(s) found**") | |
| # else: | |
| # summary_parts.append(f"✅ **Found {count} {search_terms}(s)**") | |
| # # Show detection details | |
| # for i, det in enumerate(detections[:5]): # Show first 5 | |
| # summary_parts.append(f" • Detection {i+1}: {det['score']:.3f} confidence") | |
| # if count > 5: | |
| # summary_parts.append(f" • ... and {count-5} more detections") | |
| # else: | |
| # summary_parts.append(f"❌ **No {search_terms}(s) detected**") | |
| # summary_parts.append("💡 Try lowering the threshold or using different terms") | |
| # summary_text = "\n".join(summary_parts) | |
| # return result_image, summary_text | |
| # except Exception as e: | |
| # error_msg = f"❌ **Error**: {str(e)}" | |
| # return image, error_msg | |
| # # ---------------- | |
| # # GRADIO INTERFACE | |
| # # ---------------- | |
| # with gr.Blocks(title="🔍 Object Detection & Segmentation") as demo: | |
| # gr.Markdown(""" | |
| # # 🔍 Object Detection & Segmentation App | |
| # **Simple and powerful object detection using OWL-ViT + SAM2** | |
| # 1. **Enter your query** (e.g., "How many people?", "Find cars", "Count apples") | |
| # 2. **Upload an image** | |
| # 3. **Adjust detection sensitivity** | |
| # 4. **Toggle SAM2 segmentation** for precise masks | |
| # 5. **Click Detect!** | |
| # """) | |
| # with gr.Row(): | |
| # with gr.Column(scale=1): | |
| # query_input = gr.Textbox( | |
| # label="🗣️ What do you want to detect?", | |
| # placeholder="e.g., 'How many people are in the image?'", | |
| # value="How many people are in the image?", | |
| # lines=2 | |
| # ) | |
| # image_input = gr.Image( | |
| # label="📸 Upload your image", | |
| # type="numpy" | |
| # ) | |
| # with gr.Row(): | |
| # threshold_slider = gr.Slider( | |
| # minimum=0.01, | |
| # maximum=0.9, | |
| # value=0.1, | |
| # step=0.01, | |
| # label="🎚️ Detection Sensitivity" | |
| # ) | |
| # sam_checkbox = gr.Checkbox( | |
| # label="🎭 Enable SAM2 Segmentation", | |
| # value=False, | |
| # info="Generate precise pixel masks" | |
| # ) | |
| # detect_button = gr.Button("🔍 Detect Objects!", variant="primary", size="lg") | |
| # with gr.Column(scale=1): | |
| # output_image = gr.Image(label="🎯 Detection Results") | |
| # output_text = gr.Textbox( | |
| # label="📊 Detection Summary", | |
| # lines=12, | |
| # show_copy_button=True | |
| # ) | |
| # # Event handlers | |
| # detect_button.click( | |
| # fn=detection_pipeline, | |
| # inputs=[query_input, image_input, threshold_slider, sam_checkbox], | |
| # outputs=[output_image, output_text] | |
| # ) | |
| # # Also trigger on Enter in text box | |
| # query_input.submit( | |
| # fn=detection_pipeline, | |
| # inputs=[query_input, image_input, threshold_slider, sam_checkbox], | |
| # outputs=[output_image, output_text] | |
| # ) | |
| # # Examples section | |
| # gr.Examples( | |
| # examples=[ | |
| # ["How many people are in the image?", None, 0.1, False], | |
| # ["Find all cars", None, 0.15, True], | |
| # ["Count the bottles", None, 0.1, True], | |
| # ["Detect dogs", None, 0.2, False], | |
| # ["How many phones?", None, 0.15, True], | |
| # ], | |
| # inputs=[query_input, image_input, threshold_slider, sam_checkbox], | |
| # ) | |
| # # Launch | |
| # if __name__ == "__main__": | |
| # demo.launch(server_name="0.0.0.0", server_port=7860, share=True) | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| import matplotlib.pyplot as plt | |
| from transformers import pipeline | |
| import warnings | |
| from io import BytesIO | |
| import importlib.util | |
| import os | |
| import openai | |
| from typing import List, Dict | |
| # Suppress warnings | |
| warnings.filterwarnings("ignore") | |
| # Set up OpenAI API key | |
| api_key = os.getenv('OPENAI_API_KEY') | |
| if not api_key: | |
| print("No OpenAI API key found - will use simple keyword extraction") | |
| elif not api_key.startswith("sk-proj-") and not api_key.startswith("sk-"): | |
| print("API key found but doesn't look correct") | |
| elif api_key.strip() != api_key: | |
| print("API key has leading or trailing whitespace - please fix it.") | |
| else: | |
| print("OpenAI API key found and looks good!") | |
| openai.api_key = api_key | |
| # Global variables for models | |
| detector = None | |
| sam_predictor = None | |
| def calculate_bbox_area(bbox): | |
| """Calculate the area of a normalized bounding box.""" | |
| x1, y1, x2, y2 = bbox | |
| width = abs(x2 - x1) | |
| height = abs(y2 - y1) | |
| return width * height | |
| def filter_bbox_outliers(detections: List[Dict], | |
| method: str = 'zscore', | |
| threshold: float = 2.0, | |
| min_score: float = 0.0) -> List[Dict]: | |
| """ | |
| Filter out outlier bounding boxes based on their area. | |
| Args: | |
| detections: List of detection dictionaries with 'bbox', 'label', 'score' | |
| method: 'iqr' (Interquartile Range) or 'zscore' (Z-score) | |
| threshold: Multiplier for IQR method or Z-score threshold | |
| min_score: Minimum confidence score to keep detection | |
| Returns: | |
| Filtered list of detections | |
| """ | |
| if not detections: | |
| return detections | |
| # Filter by minimum score first | |
| detections = [det for det in detections if det['score'] >= min_score] | |
| if len(detections) <= 2: # Need at least 3 detections for meaningful outlier removal | |
| return detections | |
| # Calculate areas for all bounding boxes | |
| areas = [calculate_bbox_area(det['bbox']) for det in detections] | |
| areas = np.array(areas) | |
| if method == 'iqr': | |
| # IQR method | |
| q1 = np.percentile(areas, 25) | |
| q3 = np.percentile(areas, 75) | |
| iqr = q3 - q1 | |
| lower_bound = q1 - threshold * iqr | |
| upper_bound = q3 + threshold * iqr | |
| valid_indices = np.where((areas >= lower_bound) & (areas <= upper_bound))[0] | |
| elif method == 'zscore': | |
| # Z-score method | |
| if np.std(areas) == 0: # All areas are the same | |
| return detections | |
| mean_area = np.mean(areas) | |
| std_area = np.std(areas) | |
| z_scores = np.abs((areas - mean_area) / std_area) | |
| valid_indices = np.where(z_scores <= threshold)[0] | |
| else: | |
| raise ValueError("Method must be 'iqr' or 'zscore'") | |
| # Return filtered detections | |
| filtered_detections = [detections[i] for i in valid_indices] | |
| print(f"Original detections: {len(detections)}") | |
| print(f"Filtered detections: {len(filtered_detections)}") | |
| print(f"Removed {len(detections) - len(filtered_detections)} outliers") | |
| return filtered_detections | |
| """Load the OWL-ViT detector once and cache it.""" | |
| global detector | |
| if detector is None: | |
| print("Loading OWL-ViT model...") | |
| detector = pipeline( | |
| model="google/owlv2-base-patch16-ensemble", | |
| task="zero-shot-object-detection", | |
| device=0 if torch.cuda.is_available() else -1 | |
| ) | |
| print("OWL-ViT model loaded successfully!") | |
| def install_sam2_if_needed(): | |
| """Check if SAM2 is installed, and install it if needed.""" | |
| if importlib.util.find_spec("sam2") is not None: | |
| print("SAM2 is already installed.") | |
| return True | |
| try: | |
| import subprocess | |
| import sys | |
| print("Installing SAM2 from GitHub...") | |
| subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/facebookresearch/sam2.git"]) | |
| print("SAM2 installed successfully.") | |
| return True | |
| except Exception as e: | |
| print(f"Error installing SAM2: {e}") | |
| return False | |
| def load_sam_predictor(): | |
| """Load SAM2 predictor if available.""" | |
| global sam_predictor | |
| if sam_predictor is None: | |
| if install_sam2_if_needed(): | |
| try: | |
| from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| print("Loading SAM2 model...") | |
| sam_predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| sam_predictor.model.to(device) | |
| print(f"SAM2 model loaded successfully on {device}!") | |
| return True | |
| except Exception as e: | |
| print(f"Error loading SAM2: {e}") | |
| return False | |
| return sam_predictor is not None | |
| def detect_objects_owlv2(text_query, image, threshold=0.1): | |
| """Detect objects using OWL-ViT.""" | |
| try: | |
| load_detector() | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| # Clean up the text query | |
| query_terms = [term.strip() for term in text_query.split(',') if term.strip()] | |
| if not query_terms: | |
| query_terms = ["object"] | |
| print(f"Detecting: {query_terms}") | |
| predictions = detector(image, candidate_labels=query_terms) | |
| detections = [] | |
| for pred in predictions: | |
| if pred['score'] >= threshold: | |
| bbox = pred['box'] | |
| width, height = image.size | |
| normalized_bbox = [ | |
| bbox['xmin'] / width, | |
| bbox['ymin'] / height, | |
| bbox['xmax'] / width, | |
| bbox['ymax'] / height | |
| ] | |
| detection = { | |
| 'label': pred['label'], | |
| 'bbox': normalized_bbox, | |
| 'score': pred['score'] | |
| } | |
| detections.append(detection) | |
| return detections, image | |
| except Exception as e: | |
| print(f"Detection error: {e}") | |
| return [], image | |
| def generate_masks_sam2(detections, image): | |
| """Generate segmentation masks using SAM2.""" | |
| try: | |
| if not load_sam_predictor(): | |
| print("SAM2 not available, skipping mask generation") | |
| return detections | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| image_np = np.array(image.convert("RGB")) | |
| H, W = image_np.shape[:2] | |
| # Set image for SAM2 | |
| sam_predictor.set_image(image_np) | |
| # Convert normalized bboxes to pixel coordinates | |
| input_boxes = [] | |
| for det in detections: | |
| x1, y1, x2, y2 = det['bbox'] | |
| input_boxes.append([int(x1 * W), int(y1 * H), int(x2 * W), int(y2 * H)]) | |
| if not input_boxes: | |
| return detections | |
| input_boxes = np.array(input_boxes) | |
| print(f"Generating masks for {len(input_boxes)} detections...") | |
| with torch.inference_mode(): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if device == "cuda": | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| masks, scores, _ = sam_predictor.predict( | |
| point_coords=None, | |
| point_labels=None, | |
| box=input_boxes, | |
| multimask_output=False | |
| ) | |
| else: | |
| masks, scores, _ = sam_predictor.predict( | |
| point_coords=None, | |
| point_labels=None, | |
| box=input_boxes, | |
| multimask_output=False | |
| ) | |
| # Add masks to detections | |
| results = [] | |
| for i, det in enumerate(detections): | |
| new_det = det.copy() | |
| mask = masks[i] | |
| if mask.ndim == 3: | |
| mask = mask[0] # Remove batch dimension if present | |
| new_det['mask'] = mask.astype(np.uint8) | |
| results.append(new_det) | |
| print(f"Successfully generated {len(results)} masks") | |
| return results | |
| except Exception as e: | |
| print(f"SAM2 mask generation error: {e}") | |
| return detections | |
| def visualize_detections_with_masks(image, detections_with_masks, show_labels=True, show_boxes=True): | |
| """ | |
| Visualize the detections with their segmentation masks. | |
| Returns PIL Image instead of showing plot. | |
| """ | |
| # Load the image | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| image_np = np.array(image.convert("RGB")) | |
| # Get image dimensions | |
| height, width = image_np.shape[:2] | |
| # Create figure | |
| fig = plt.figure(figsize=(12, 8)) | |
| plt.imshow(image_np) | |
| # Define colors for different instances | |
| colors = plt.cm.tab10(np.linspace(0, 1, 10)) | |
| # Plot each detection | |
| for i, detection in enumerate(detections_with_masks): | |
| # Get bbox, mask, label, and score | |
| bbox = detection['bbox'] | |
| label = detection['label'] | |
| score = detection['score'] | |
| # Convert normalized bbox to pixel coordinates | |
| x1, y1, x2, y2 = bbox | |
| x1_px, y1_px = int(x1 * width), int(y1 * height) | |
| x2_px, y2_px = int(x2 * width), int(y2 * height) | |
| # Color for this instance | |
| color = colors[i % len(colors)] | |
| # Display mask with transparency if available | |
| if 'mask' in detection: | |
| mask = detection['mask'] | |
| mask_color = np.zeros((height, width, 4), dtype=np.float32) | |
| mask_color[mask > 0] = [color[0], color[1], color[2], 0.5] | |
| plt.imshow(mask_color) | |
| # Draw bounding box if requested | |
| if show_boxes: | |
| rect = plt.Rectangle((x1_px, y1_px), x2_px - x1_px, y2_px - y1_px, | |
| fill=False, edgecolor=color, linewidth=2) | |
| plt.gca().add_patch(rect) | |
| # Add label and score if requested | |
| if show_labels: | |
| plt.text(x1_px, y1_px - 5, f"{label}: {score:.2f}", | |
| color='white', bbox=dict(facecolor=color, alpha=0.8), fontsize=10) | |
| plt.axis('off') | |
| plt.tight_layout() | |
| # Convert to PIL Image using the correct method | |
| buf = BytesIO() | |
| plt.savefig(buf, format='png', bbox_inches='tight', dpi=150) | |
| plt.close(fig) | |
| buf.seek(0) | |
| result_image = Image.open(buf) | |
| return result_image | |
| def visualize_detections(image, detections, show_labels=True): | |
| """ | |
| Visualize object detections with bounding boxes only. | |
| Returns PIL Image instead of showing plot. | |
| """ | |
| # Load the image | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| image_np = np.array(image.convert("RGB")) | |
| # Get image dimensions | |
| height, width = image_np.shape[:2] | |
| # Create figure | |
| fig = plt.figure(figsize=(12, 8)) | |
| plt.imshow(image_np) | |
| # If we have detections, draw them | |
| if detections: | |
| # Define colors for different instances | |
| colors = plt.cm.tab10(np.linspace(0, 1, 10)) | |
| # Plot each detection | |
| for i, detection in enumerate(detections): | |
| # Get bbox, label, and score | |
| bbox = detection['bbox'] | |
| label = detection['label'] | |
| score = detection['score'] | |
| # Convert normalized bbox to pixel coordinates | |
| x1, y1, x2, y2 = bbox | |
| x1_px, y1_px = int(x1 * width), int(y1 * height) | |
| x2_px, y2_px = int(x2 * width), int(y2 * height) | |
| # Color for this instance | |
| color = colors[i % len(colors)] | |
| # Draw bounding box | |
| rect = plt.Rectangle((x1_px, y1_px), x2_px - x1_px, y2_px - y1_px, | |
| fill=False, edgecolor=color, linewidth=2) | |
| plt.gca().add_patch(rect) | |
| # Add label and score if requested | |
| if show_labels: | |
| plt.text(x1_px, y1_px - 5, f"{label}: {score:.2f}", | |
| color='white', bbox=dict(facecolor=color, alpha=0.8), fontsize=10) | |
| # Set title | |
| plt.title(f'Object Detection Results ({len(detections)} objects found)', fontsize=14, pad=20) | |
| plt.axis('off') | |
| plt.tight_layout() | |
| # Convert to PIL Image | |
| buf = BytesIO() | |
| plt.savefig(buf, format='png', bbox_inches='tight', dpi=150) | |
| plt.close(fig) | |
| buf.seek(0) | |
| result_image = Image.open(buf) | |
| return result_image | |
| def get_optimized_prompt(query_text): | |
| """ | |
| Use OpenAI to convert natural language query into optimal detection prompt. | |
| Falls back to simple extraction if OpenAI is not available. | |
| """ | |
| if not query_text.strip(): | |
| return "object" | |
| # Try OpenAI first if API key is available | |
| if hasattr(openai, 'api_key') and openai.api_key: | |
| try: | |
| response = openai.chat.completions.create( | |
| model="gpt-3.5-turbo", | |
| messages=[{ | |
| "role": "system", | |
| "content": """You are an expert at converting natural language queries into precise object detection terms. | |
| RULES: | |
| 1. Return ONLY 1-2 words maximum that describe the object to detect | |
| 2. Use the exact object name from the user's query | |
| 3. For people: use "person" | |
| 4. For vehicles: use "car", "truck", "bicycle" | |
| 5. Do NOT include counting words, articles, or explanations | |
| 6. Examples: | |
| - "How many cacao fruits are there?" → "cacao fruit" | |
| - "Count the corn in the field" → "corn" | |
| - "Find all people" → "person" | |
| - "How many cacao pods?" → "cacao pod" | |
| - "Detect cars" → "car" | |
| - "Count bananas" → "banana" | |
| - "How many apples?" → "apple" | |
| Return ONLY the object name, nothing else.""" | |
| }, { | |
| "role": "user", | |
| "content": query_text | |
| }], | |
| temperature=0.0, # Make it deterministic | |
| max_tokens=5 # Force brevity | |
| ) | |
| llm_result = response.choices[0].message.content.strip().lower() | |
| # Extra safety: take only first 2 words | |
| words = llm_result.split()[:2] | |
| final_result = " ".join(words) | |
| print(f"🤖 OpenAI suggested prompt: '{final_result}'") | |
| return final_result | |
| except Exception as e: | |
| print(f"OpenAI error: {e}, falling back to keyword extraction") | |
| # Fallback to simple keyword extraction (no hardcoded fruits) | |
| print("🔤 Using keyword extraction (no OpenAI)") | |
| query_lower = query_text.lower().replace("?", "").strip() | |
| # Look for common patterns and extract object names | |
| if "how many" in query_lower: | |
| parts = query_lower.split("how many") | |
| if len(parts) > 1: | |
| remaining = parts[1].strip() | |
| remaining = remaining.replace("are", "").replace("in", "").replace("the", "").replace("image", "").replace("there", "").strip() | |
| # Take first meaningful word(s) | |
| words = remaining.split()[:2] | |
| search_terms = " ".join(words) if words else "object" | |
| else: | |
| search_terms = "object" | |
| elif "count" in query_lower: | |
| parts = query_lower.split("count") | |
| if len(parts) > 1: | |
| remaining = parts[1].strip() | |
| remaining = remaining.replace("the", "").replace("in", "").replace("image", "").strip() | |
| words = remaining.split()[:2] | |
| search_terms = " ".join(words) if words else "object" | |
| else: | |
| search_terms = "object" | |
| elif "find" in query_lower: | |
| parts = query_lower.split("find") | |
| if len(parts) > 1: | |
| remaining = parts[1].strip() | |
| remaining = remaining.replace("all", "").replace("the", "").replace("in", "").replace("image", "").strip() | |
| words = remaining.split()[:2] | |
| search_terms = " ".join(words) if words else "object" | |
| else: | |
| search_terms = "object" | |
| else: | |
| # Extract first 1-2 meaningful words from the query | |
| words = query_lower.split() | |
| meaningful_words = [w for w in words if w not in ["how", "many", "are", "in", "the", "image", "find", "count", "detect", "there", "this", "that", "a", "an"]] | |
| search_terms = " ".join(meaningful_words[:2]) if meaningful_words else "object" | |
| return search_terms | |
| def is_count_query(text): | |
| """Check if the query is asking for counting.""" | |
| count_keywords = ["how many", "count", "number of", "total"] | |
| return any(keyword in text.lower() for keyword in count_keywords) | |
| def detection_pipeline(query_text, image, threshold, use_sam): | |
| """Main detection pipeline.""" | |
| if image is None: | |
| return None, "⚠️ Please upload an image first!" | |
| try: | |
| # Use OpenAI or fallback to get optimized search terms | |
| search_terms = get_optimized_prompt(query_text) | |
| print(f"Processing query: '{query_text}' -> searching for: '{search_terms}'") | |
| # Run object detection | |
| detections, processed_image = detect_objects_owlv2(search_terms, image, threshold) | |
| print(f"Found {len(detections)} initial detections") | |
| for i, det in enumerate(detections): | |
| print(f"Detection {i+1}: {det['label']} (score: {det['score']:.3f}, area: {calculate_bbox_area(det['bbox']):.6f})") | |
| # Filter outliers before SAM2 | |
| if len(detections) > 2: # Only filter if we have enough detections | |
| detections = filter_bbox_outliers(detections, method='zscore', threshold=2.0) | |
| print(f"After outlier filtering: {len(detections)} detections remain") | |
| # Generate masks if requested | |
| if use_sam and detections: | |
| print("Generating SAM2 masks...") | |
| detections = generate_masks_sam2(detections, processed_image) | |
| # Create visualization using your proven functions (labels OFF) | |
| print("Creating visualization...") | |
| if use_sam and detections and 'mask' in detections[0]: | |
| result_image = visualize_detections_with_masks( | |
| processed_image, | |
| detections, | |
| show_labels=False, # Labels OFF | |
| show_boxes=True | |
| ) | |
| print("Created visualization with masks") | |
| else: | |
| result_image = visualize_detections( | |
| processed_image, | |
| detections, | |
| show_labels=False # Labels OFF | |
| ) | |
| print("Created visualization with bounding boxes only") | |
| # Make sure we have a valid result image | |
| if result_image is None: | |
| print("Warning: result_image is None, returning original image") | |
| result_image = processed_image | |
| # Generate summary | |
| count = len(detections) | |
| summary_parts = [] | |
| summary_parts.append(f"🗣️ **Original Query**: '{query_text}'") | |
| summary_parts.append(f"🤖 **AI-Optimized Search**: '{search_terms}'") | |
| summary_parts.append(f"⚙️ **Threshold**: {threshold}") | |
| summary_parts.append(f"🎭 **SAM2 Segmentation**: {'Enabled' if use_sam else 'Disabled'}") | |
| if count > 0: | |
| if is_count_query(query_text): | |
| summary_parts.append(f"🔢 **Answer: {count} {search_terms}(s) found**") | |
| else: | |
| summary_parts.append(f"✅ **Found {count} {search_terms}(s)**") | |
| # Show detection details | |
| for i, det in enumerate(detections[:5]): # Show first 5 | |
| summary_parts.append(f" • Detection {i+1}: {det['score']:.3f} confidence") | |
| if count > 5: | |
| summary_parts.append(f" • ... and {count-5} more detections") | |
| else: | |
| summary_parts.append(f"❌ **No {search_terms}(s) detected**") | |
| summary_parts.append("💡 Try lowering the threshold or using different terms") | |
| summary_text = "\n".join(summary_parts) | |
| return result_image, summary_text | |
| except Exception as e: | |
| error_msg = f"❌ **Error**: {str(e)}" | |
| return image, error_msg | |
| # ---------------- | |
| # GRADIO INTERFACE | |
| # ---------------- | |
| with gr.Blocks(title="🔍 Object Detection & Segmentation") as demo: | |
| gr.Markdown(""" | |
| # 🔍 Object Detection & Segmentation App | |
| **Simple and powerful object detection using OWL-ViT + SAM2** | |
| 1. **Enter your query** (e.g., "How many people?", "Find cars", "Count apples") | |
| 2. **Upload an image** | |
| 3. **Adjust detection sensitivity** | |
| 4. **Toggle SAM2 segmentation** for precise masks | |
| 5. **Click Detect!** | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| query_input = gr.Textbox( | |
| label="🗣️ What do you want to detect?", | |
| placeholder="e.g., 'How many people are in the image?'", | |
| value="How many people are in the image?", | |
| lines=2 | |
| ) | |
| image_input = gr.Image( | |
| label="📸 Upload your image", | |
| type="numpy" | |
| ) | |
| with gr.Row(): | |
| threshold_slider = gr.Slider( | |
| minimum=0.01, | |
| maximum=0.9, | |
| value=0.1, | |
| step=0.01, | |
| label="🎚️ Detection Sensitivity" | |
| ) | |
| sam_checkbox = gr.Checkbox( | |
| label="🎭 Enable SAM2 Segmentation", | |
| value=False, | |
| info="Generate precise pixel masks" | |
| ) | |
| detect_button = gr.Button("🔍 Detect Objects!", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| output_image = gr.Image(label="🎯 Detection Results") | |
| output_text = gr.Textbox( | |
| label="📊 Detection Summary", | |
| lines=12, | |
| show_copy_button=True | |
| ) | |
| # Event handlers | |
| detect_button.click( | |
| fn=detection_pipeline, | |
| inputs=[query_input, image_input, threshold_slider, sam_checkbox], | |
| outputs=[output_image, output_text] | |
| ) | |
| # Also trigger on Enter in text box | |
| query_input.submit( | |
| fn=detection_pipeline, | |
| inputs=[query_input, image_input, threshold_slider, sam_checkbox], | |
| outputs=[output_image, output_text] | |
| ) | |
| # Examples section | |
| gr.Examples( | |
| examples=[ | |
| ["How many people are in the image?", None, 0.1, False], | |
| ["Find all cars", None, 0.15, True], | |
| ["Count the bottles", None, 0.1, True], | |
| ["Detect dogs", None, 0.2, False], | |
| ["How many phones?", None, 0.15, True], | |
| ], | |
| inputs=[query_input, image_input, threshold_slider, sam_checkbox], | |
| ) | |
| # Launch | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, share=True) |