import gradio as gr import torch from PIL import Image, ImageDraw, ImageFont from transformers import DetrImageProcessor, DetrForObjectDetection import transformers # Global variables to cache models current_model = None current_processor = None current_model_name = None # Available models with better selection available_models = { # DETR Models "DETR ResNet-50": "facebook/detr-resnet-50", "DETR ResNet-101": "facebook/detr-resnet-101", "DETR DC5": "facebook/detr-resnet-50-dc5", "DETR ResNet-50 Face Only": "esraakh/detr_fine_tune_face_detection_final" } def load_model(model_key): """Load model and processor based on selected model key""" global current_model, current_processor, current_model_name model_name = available_models[model_key] # Only load if it's a different model if current_model_name != model_name: print(f"Loading model: {model_name}") current_processor = DetrImageProcessor.from_pretrained(model_name) current_model = DetrForObjectDetection.from_pretrained(model_name) current_model_name = model_name print(f"Model loaded: {model_name}") print(f"Available labels: {list(current_model.config.id2label.values())}") return current_model, current_processor # Fixed font loading - this was the main issue def get_font(size=12): try: return ImageFont.truetype("arial.ttf", size=size) except: return ImageFont.load_default() # Set up translations for the app translations = { "English": { "title": "## Enhanced Object Detection App\nUpload an image to detect objects using various DETR models.", "input_label": "Input Image", "output_label": "Detected Objects", "dropdown_label": "Label Language", "dropdown_detection_model_label": "Detection Model", "threshold_label": "Detection Threshold", "button": "Detect Objects", "info_label": "Detection Info", "model_fast": "General Objects (fast)", "model_precision": "General Objects (high precision)", "model_small": "Small Objects/Details (slow)", "model_faces": "Face Detection (people only)" }, "Spanish": { "title": "## Aplicación Mejorada de Detección de Objetos\nSube una imagen para detectar objetos usando varios modelos DETR.", "input_label": "Imagen de entrada", "output_label": "Objetos detectados", "dropdown_label": "Idioma de las etiquetas", "dropdown_detection_model_label": "Modelo de detección", "threshold_label": "Umbral de detección", "button": "Detectar objetos", "info_label": "Información de detección", "model_fast": "Objetos generales (rápido)", "model_precision": "Objetos generales (precisión alta)", "model_small": "Objetos pequeños/detalles (lento)", "model_faces": "Detección de caras (solo personas)" }, "French": { "title": "## Application Améliorée de Détection d'Objets\nTéléchargez une image pour détecter des objets avec divers modèles DETR.", "input_label": "Image d'entrée", "output_label": "Objets détectés", "dropdown_label": "Langue des étiquettes", "dropdown_detection_model_label": "Modèle de détection", "threshold_label": "Seuil de détection", "button": "Détecter les objets", "info_label": "Information de détection", "model_fast": "Objets généraux (rapide)", "model_precision": "Objets généraux (haute précision)", "model_small": "Petits objets/détails (lent)", "model_faces": "Détection de visages (personnes uniquement)" } } def t(language, key): return translations.get(language, translations["English"]).get(key, key) def get_translated_model_choices(language): """Get model choices translated to the selected language""" model_mapping = { "DETR ResNet-50": "model_fast", "DETR ResNet-101": "model_precision", "DETR DC5": "model_small", "DETR ResNet-50 Face Only": "model_faces" } translated_choices = [] for model_key in available_models.keys(): if model_key in model_mapping: translation_key = model_mapping[model_key] translated_name = t(language, translation_key) else: translated_name = model_key # Fallback to original name translated_choices.append(translated_name) return translated_choices def get_model_key_from_translation(translated_name, language): """Get the original model key from translated name""" model_mapping = { "DETR ResNet-50": "model_fast", "DETR ResNet-101": "model_precision", "DETR DC5": "model_small", "DETR ResNet-50 Face Only": "model_faces" } # Reverse lookup for model_key, translation_key in model_mapping.items(): if t(language, translation_key) == translated_name: return model_key # If not found, try direct match if translated_name in available_models: return translated_name # Default fallback return "DETR ResNet-50" def get_helsinki_model(language_label): """Returns the Helsinki-NLP model name for translating from English to the selected language.""" lang_map = { "Spanish": "es", "French": "fr", "English": "en" } target = lang_map.get(language_label) if not target or target == "en": return None return f"Helsinki-NLP/opus-mt-en-{target}" # add cache for translations translation_cache = {} def translate_label(language_label, label): """Translates the given label to the target language.""" # Check cache first cache_key = f"{language_label}_{label}" if cache_key in translation_cache: return translation_cache[cache_key] model_name = get_helsinki_model(language_label) if not model_name: return label try: translator = transformers.pipeline("translation", model=model_name) result = translator(label, max_length=40) translated = result[0]['translation_text'] # Cache the result translation_cache[cache_key] = translated return translated except Exception as e: print(f"Translation error (429 or other): {e}") return label # Return original if translation fails def detect_objects(image, language_selector, translated_model_selector, threshold): """Enhanced object detection with adjustable threshold and better info""" try: if image is None: return None, "Por favor, sube una imagen antes de detectar objetos." model_selector = get_model_key_from_translation(translated_model_selector, language_selector) print(f"Processing image. Language: {language_selector}, Model: {model_selector}, Threshold: {threshold}") model, processor = load_model(model_selector) inputs = processor(images=image, return_tensors="pt") outputs = model(**inputs) target_sizes = torch.tensor([image.size[::-1]]) results = processor.post_process_object_detection( outputs, threshold=threshold, target_sizes=target_sizes )[0] image_with_boxes = image.copy() draw = ImageDraw.Draw(image_with_boxes) detection_info = f"Detected {len(results['scores'])} objects with threshold {threshold}\n" detection_info += f"Model: {translated_model_selector} ({model_selector})\n\n" colors = { 'high': 'red', # > 0.8 'medium': 'orange', # 0.5-0.8 'low': 'yellow' # < 0.5 } detected_objects = [] for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): confidence = score.item() box = [round(x, 2) for x in box.tolist()] if confidence > 0.8: color = colors['high'] elif confidence > 0.5: color = colors['medium'] else: color = colors['low'] draw.rectangle(box, outline=color, width=3) label_text = model.config.id2label[label.item()] translated_label = translate_label(language_selector, label_text) display_text = f"{translated_label}: {round(confidence, 3)}" detected_objects.append({ 'label': label_text, 'translated': translated_label, 'confidence': confidence, 'box': box }) try: image_width = image.size[0] font_size = max(image_width // 40, 12) font = get_font(font_size) text_bbox = draw.textbbox((0, 0), display_text, font=font) text_width = text_bbox[2] - text_bbox[0] text_height = text_bbox[3] - text_bbox[1] except: font = get_font(12) text_width = 50 text_height = 20 text_bg = [ box[0], box[1] - text_height - 4, box[0] + text_width + 4, box[1] ] draw.rectangle(text_bg, fill="black") draw.text((box[0] + 2, box[1] - text_height - 2), display_text, fill="white", font=font) if detected_objects: detection_info += "Objects found:\n" for obj in sorted(detected_objects, key=lambda x: x['confidence'], reverse=True): detection_info += f"- {obj['translated']} ({obj['label']}): {obj['confidence']:.3f}\n" else: detection_info += "No objects detected. Try lowering the threshold." return image_with_boxes, detection_info except Exception as e: import traceback print("ERROR EN DETECT_OBJECTS:", e) traceback.print_exc() return None, f"Error detecting objects: {e}" def build_app(): with gr.Blocks(theme=gr.themes.Soft()) as app: with gr.Row(): title = gr.Markdown(t("English", "title")) with gr.Row(): with gr.Column(scale=1): language_selector = gr.Dropdown( choices=["English", "Spanish", "French"], value="English", label=t("English", "dropdown_label") ) with gr.Column(scale=1): model_selector = gr.Dropdown( choices=get_translated_model_choices("English"), value=t("English", "model_fast"), # Default to translated "fast" option label=t("English", "dropdown_detection_model_label") ) with gr.Column(scale=1): threshold_slider = gr.Slider( minimum=0.1, maximum=0.95, value=0.5, # Lowered default threshold step=0.05, label=t("English", "threshold_label") ) with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(type="pil", label=t("English", "input_label")) button = gr.Button(t("English", "button"), variant="primary") with gr.Column(scale=1): output_image = gr.Image(label=t("English", "output_label")) detection_info = gr.Textbox( label=t("English", "info_label"), lines=10, max_lines=15 ) # Function to update interface when language changes def update_interface(selected_language): translated_choices = get_translated_model_choices(selected_language) default_model = t(selected_language, "model_fast") return [ gr.update(value=t(selected_language, "title")), gr.update(label=t(selected_language, "dropdown_label")), gr.update( choices=translated_choices, value=default_model, label=t(selected_language, "dropdown_detection_model_label") ), gr.update(label=t(selected_language, "threshold_label")), gr.update(label=t(selected_language, "input_label")), gr.update(value=t(selected_language, "button")), gr.update(label=t(selected_language, "output_label")), gr.update(label=t(selected_language, "info_label")) ] # Connect language change event language_selector.change( fn=update_interface, inputs=language_selector, outputs=[title, language_selector, model_selector, threshold_slider, input_image, button, output_image, detection_info], queue=False ) # Connect detection button click event button.click( fn=detect_objects, inputs=[input_image, language_selector, model_selector, threshold_slider], outputs=[output_image, detection_info] ) return app # Initialize with default model load_model("DETR ResNet-50") # Launch the application if __name__ == "__main__": app = build_app() app.launch()