aggiunte immagini

gradio_ui/app.py

@@ -6,6 +6,8 @@ by assembling all tab components and wiring event handlers.
 """
 
 import gradio as gr
+import numpy as np
+from PIL import Image
 
 from core import get_available_models, MVTEC_CATEGORIES
 from gradio_ui.handlers import (
@@ -15,6 +17,7 @@ from gradio_ui.handlers import (
     predict_compare,
     update_categories,
     update_compare_categories,
+    get_category_from_sample,
 )
 from gradio_ui.content import MAIN_HEADER
 from gradio_ui.tabs import (
@@ -69,7 +72,7 @@ def create_app() -> gr.Blocks:
         # =========================================================================
         upload_components["model_dropdown"].change(
             fn=update_categories,
-            inputs=[upload_components["model_dropdown"]],
+            inputs=[upload_components["model_dropdown"], upload_components["category_dropdown"]],
             outputs=[upload_components["category_dropdown"]]
         )
         
@@ -86,18 +89,42 @@ def create_app() -> gr.Blocks:
             ]
         )
         
-        # Allow Enter key on image upload
-        upload_components["image_input"].change(
-            fn=predict,
-            inputs=[
-                upload_components["image_input"],
-                upload_components["model_dropdown"],
-                upload_components["category_dropdown"]
-            ],
-            outputs=[
-                upload_components["result_image"],
-                upload_components["result_text"]
-            ]
+        # Sample gallery selection handler for Upload tab
+        def on_upload_sample_select(evt: gr.SelectData, current_model: str, gallery_value: list):
+            """Handler for sample image selection in upload tab."""
+            if evt is None:
+                return None, gr.Dropdown()
+            
+            # Get the selected index and retrieve original path from gallery data
+            idx = evt.index
+            if gallery_value and idx < len(gallery_value):
+                item = gallery_value[idx]
+                if isinstance(item, tuple):
+                    image_path, label = item
+                    category = get_category_from_sample(label)
+                else:
+                    image_path = item
+                    category = get_category_from_sample(str(image_path))
+            else:
+                return None, gr.Dropdown()
+            
+            try:
+                img = np.array(Image.open(image_path).convert("RGB"))
+            except Exception as e:
+                print(f"Error loading image: {e}")
+                return None, gr.Dropdown()
+            
+            trained = get_trained_categories(current_model)
+            if trained and category in trained:
+                return img, gr.Dropdown(choices=trained, value=category)
+            elif trained:
+                return img, gr.Dropdown(choices=trained, value=trained[0])
+            return img, gr.Dropdown(choices=MVTEC_CATEGORIES, value=category)
+        
+        upload_components["sample_gallery"].select(
+            fn=on_upload_sample_select,
+            inputs=[upload_components["model_dropdown"], upload_components["sample_gallery"]],
+            outputs=[upload_components["image_input"], upload_components["category_dropdown"]]
         )
         
         # =========================================================================
@@ -105,7 +132,7 @@ def create_app() -> gr.Blocks:
         # =========================================================================
         sketch_components["model_dropdown"].change(
             fn=update_categories,
-            inputs=[sketch_components["model_dropdown"]],
+            inputs=[sketch_components["model_dropdown"], sketch_components["category_dropdown"]],
             outputs=[sketch_components["category_dropdown"]]
         )
         
@@ -122,12 +149,52 @@ def create_app() -> gr.Blocks:
             ]
         )
         
+        # Sample gallery selection handler for Sketch tab
+        def on_sketch_sample_select(evt: gr.SelectData, current_model: str, gallery_value: list):
+            """Handler for sample image selection in sketch tab."""
+            if evt is None:
+                return None, gr.Dropdown()
+            
+            # Get the selected index and retrieve original path from gallery data
+            idx = evt.index
+            if gallery_value and idx < len(gallery_value):
+                item = gallery_value[idx]
+                if isinstance(item, tuple):
+                    image_path, label = item
+                    category = get_category_from_sample(label)
+                else:
+                    image_path = item
+                    category = get_category_from_sample(str(image_path))
+            else:
+                return None, gr.Dropdown()
+            
+            try:
+                img = np.array(Image.open(image_path).convert("RGB"))
+                # For ImageEditor, return the image in the expected format
+                editor_value = {"background": img, "layers": [], "composite": img}
+            except Exception as e:
+                print(f"Error loading image: {e}")
+                return None, gr.Dropdown()
+            
+            trained = get_trained_categories(current_model)
+            if trained and category in trained:
+                return editor_value, gr.Dropdown(choices=trained, value=category)
+            elif trained:
+                return editor_value, gr.Dropdown(choices=trained, value=trained[0])
+            return editor_value, gr.Dropdown(choices=MVTEC_CATEGORIES, value=category)
+        
+        sketch_components["sample_gallery"].select(
+            fn=on_sketch_sample_select,
+            inputs=[sketch_components["model_dropdown"], sketch_components["sample_gallery"]],
+            outputs=[sketch_components["image_editor"], sketch_components["category_dropdown"]]
+        )
+        
         # =========================================================================
         # Event handlers for Tab 3 (Compare)
         # =========================================================================
         compare_components["model_checkboxes"].change(
             fn=update_compare_categories,
-            inputs=[compare_components["model_checkboxes"]],
+            inputs=[compare_components["model_checkboxes"], compare_components["category_dropdown"]],
             outputs=[compare_components["category_dropdown"]]
         )
         
@@ -143,6 +210,46 @@ def create_app() -> gr.Blocks:
                 compare_components["summary"]
             ]
         )
+        
+        # Sample gallery selection handler for Compare tab
+        def on_compare_sample_select(evt: gr.SelectData, selected_models: list, gallery_value: list):
+            """Handler for sample image selection in compare tab."""
+            if evt is None:
+                return None, gr.Dropdown()
+            
+            # Get the selected index and retrieve original path from gallery data
+            idx = evt.index
+            if gallery_value and idx < len(gallery_value):
+                item = gallery_value[idx]
+                if isinstance(item, tuple):
+                    image_path, label = item
+                    category = get_category_from_sample(label)
+                else:
+                    image_path = item
+                    category = get_category_from_sample(str(image_path))
+            else:
+                return None, gr.Dropdown()
+            
+            try:
+                img = np.array(Image.open(image_path).convert("RGB"))
+            except Exception as e:
+                print(f"Error loading image: {e}")
+                return None, gr.Dropdown()
+            
+            # Get common categories for selected models
+            from gradio_ui.handlers import get_common_categories
+            common = get_common_categories(selected_models)
+            if common and category in common:
+                return img, gr.Dropdown(choices=common, value=category)
+            elif common:
+                return img, gr.Dropdown(choices=common, value=common[0])
+            return img, gr.Dropdown(choices=MVTEC_CATEGORIES, value=category)
+        
+        compare_components["sample_gallery"].select(
+            fn=on_compare_sample_select,
+            inputs=[compare_components["model_checkboxes"], compare_components["sample_gallery"]],
+            outputs=[compare_components["image_input"], compare_components["category_dropdown"]]
+        )
     
     return app
 

gradio_ui/handlers.py

@@ -256,33 +256,150 @@ def get_common_categories(model_names: list) -> list:
     return list(common) if common else MVTEC_CATEGORIES
 
 
-def update_categories(model_name: str):
+def update_categories(model_name: str, current_category: str = None):
     """
     Updates available categories based on selected model.
+    Keeps the current category if it's available for the new model.
     
     Args:
         model_name: Selected model name
+        current_category: Currently selected category to preserve if possible
     
     Returns:
         Updated Gradio Dropdown component
     """
     trained = get_trained_categories(model_name)
     if trained:
+        # Keep current category if it's available in the new model
+        if current_category and current_category in trained:
+            return gr.Dropdown(choices=trained, value=current_category)
         return gr.Dropdown(choices=trained, value=trained[0])
+    # Keep current category if using all categories
+    if current_category and current_category in MVTEC_CATEGORIES:
+        return gr.Dropdown(choices=MVTEC_CATEGORIES, value=current_category)
     return gr.Dropdown(choices=MVTEC_CATEGORIES, value="bottle")
 
 
-def update_compare_categories(selected_models: list):
+def update_compare_categories(selected_models: list, current_category: str = None):
     """
     Update categories dropdown based on selected models.
+    Keeps the current category if it's available for all selected models.
     
     Args:
         selected_models: List of selected model names
+        current_category: Currently selected category to preserve if possible
     
     Returns:
         Updated Gradio Dropdown component
     """
     common = get_common_categories(selected_models)
     if common:
+        # Keep current category if it's in common categories
+        if current_category and current_category in common:
+            return gr.Dropdown(choices=common, value=current_category)
         return gr.Dropdown(choices=common, value=common[0])
+    # Keep current category if using all categories
+    if current_category and current_category in MVTEC_CATEGORIES:
+        return gr.Dropdown(choices=MVTEC_CATEGORIES, value=current_category)
     return gr.Dropdown(choices=MVTEC_CATEGORIES, value="bottle")
+
+
+def get_sample_images() -> list:
+    """
+    Returns a list of sample images from the MVTecAD test set.
+    
+    Returns:
+        List of tuples (image_path, label) for gallery display
+    """
+    from core import DIR_DATASET
+    
+    samples = []
+    
+    for category in MVTEC_CATEGORIES:
+        test_dir = DIR_DATASET / category / "test"
+        if not test_dir.exists():
+            continue
+        
+        # Get all defect types for this category
+        for defect_type in test_dir.iterdir():
+            if not defect_type.is_dir():
+                continue
+            
+            # Get first image from each defect type
+            images = list(defect_type.glob("*.png"))
+            if images:
+                img_path = str(images[0])
+                label = f"{category}/{defect_type.name}"
+                samples.append((img_path, label))
+    
+    return samples
+
+
+def get_category_from_sample(image_path_or_label: str) -> str:
+    """
+    Extracts category from a sample image path or gallery label.
+    
+    Args:
+        image_path_or_label: Path to the sample image or gallery label (format: category/defect_type)
+    
+    Returns:
+        Category name extracted from the path or label
+    """
+    if image_path_or_label is None:
+        return "bottle"
+    
+    path_str = str(image_path_or_label)
+    
+    # First try: check if it's a label format "category/defect_type"
+    for category in MVTEC_CATEGORIES:
+        if path_str.startswith(f"{category}/") or path_str == category:
+            return category
+    
+    # Second try: check if category is in the path
+    for category in MVTEC_CATEGORIES:
+        if f"/{category}/" in path_str or f"\\{category}\\" in path_str:
+            return category
+    
+    return "bottle"
+
+
+def on_sample_select(evt: gr.SelectData, current_model: str = None):
+    """
+    Handler for when a sample image is selected from the gallery.
+    Returns the image and updates the category.
+    
+    Args:
+        evt: Gradio SelectData event with image info
+        current_model: Currently selected model (for category validation)
+    
+    Returns:
+        Tuple of (image, category_dropdown_update)
+    """
+    if evt is None or evt.value is None:
+        return None, gr.Dropdown()
+    
+    # Get the image path from the event
+    image_data = evt.value
+    if isinstance(image_data, dict):
+        image_path = image_data.get("image", {}).get("path", "")
+    else:
+        image_path = str(image_data)
+    
+    # Extract category from path
+    category = get_category_from_sample(image_path)
+    
+    # Load the image
+    try:
+        img = np.array(Image.open(image_path).convert("RGB"))
+    except Exception:
+        return None, gr.Dropdown()
+    
+    # Get available categories for current model
+    if current_model:
+        trained = get_trained_categories(current_model)
+        if trained and category in trained:
+            return img, gr.Dropdown(choices=trained, value=category)
+        elif trained:
+            return img, gr.Dropdown(choices=trained, value=trained[0])
+    
+    return img, gr.Dropdown(choices=MVTEC_CATEGORIES, value=category)

gradio_ui/tabs/compare.py

@@ -7,6 +7,7 @@ Provides the side-by-side model comparison interface.
 import gradio as gr
 
 from gradio_ui.content import COMPARE_MODELS_INSTRUCTIONS
+from gradio_ui.handlers import get_sample_images
 
 
 def create_compare_tab(available_models: list, initial_categories: list):
@@ -67,6 +68,19 @@ def create_compare_tab(available_models: list, initial_categories: list):
                 )
                 
                 compare_summary = gr.Markdown("*Select models and upload an image to compare*")
+        
+        # Sample Images Gallery
+        gr.Markdown("### 🖼️ Sample Images (click to select)")
+        compare_sample_gallery = gr.Gallery(
+            value=get_sample_images(),
+            label="Sample Images",
+            show_label=False,
+            columns=8,
+            rows=3,
+            height=350,
+            object_fit="cover",
+            allow_preview=False,
+        )
     
     return {
         "image_input": compare_image_input,
@@ -75,4 +89,5 @@ def create_compare_tab(available_models: list, initial_categories: list):
         "compare_btn": compare_btn,
         "result_image": compare_result_image,
         "summary": compare_summary,
+        "sample_gallery": compare_sample_gallery,
     }

gradio_ui/tabs/draw_defects.py

@@ -11,6 +11,7 @@ from gradio_ui.content import (
     BRUSH_COLORS_INFO,
     HEATMAP_INTERPRETATION,
 )
+from gradio_ui.handlers import get_sample_images
 
 
 def create_draw_defects_tab(available_models: list, default_model: str, initial_categories: list):
@@ -82,6 +83,19 @@ def create_draw_defects_tab(available_models: list, default_model: str, initial_
                 sketch_result_text = gr.Markdown(value="", label="Results")
                 
                 gr.Markdown(HEATMAP_INTERPRETATION)
+        
+        # Sample Images Gallery
+        gr.Markdown("### 🖼️ Sample Images (click to load into editor)")
+        sketch_sample_gallery = gr.Gallery(
+            value=get_sample_images(),
+            label="Sample Images",
+            show_label=False,
+            columns=8,
+            rows=3,
+            height=350,
+            object_fit="cover",
+            allow_preview=False,
+        )
     
     return {
         "image_editor": image_editor,
@@ -91,4 +105,5 @@ def create_draw_defects_tab(available_models: list, default_model: str, initial_
         "result_image": sketch_result_image,
         "result_image": sketch_result_image,
         "result_text": sketch_result_text,
+        "sample_gallery": sketch_sample_gallery,
     }

gradio_ui/tabs/upload.py

@@ -5,6 +5,7 @@ Provides the standard image upload and analysis interface.
 """
 
 import gradio as gr
+from gradio_ui.handlers import get_sample_images
 
 
 def create_upload_tab(available_models: list, default_model: str, initial_categories: list):
@@ -61,7 +62,20 @@ def create_upload_tab(available_models: list, default_model: str, initial_catego
                     height=300
                 )
 
-                result_text = gr.Markdown(value="", label="Results") 
+                result_text = gr.Markdown(value="", label="Results")
+        
+        # Sample Images Gallery
+        gr.Markdown("### 🖼️ Sample Images (click to select)")
+        sample_gallery = gr.Gallery(
+            value=get_sample_images(),
+            label="Sample Images",
+            show_label=False,
+            columns=8,
+            rows=3,
+            height=350,
+            object_fit="cover",
+            allow_preview=False,
+        )
         
     
     return {
@@ -72,4 +86,5 @@ def create_upload_tab(available_models: list, default_model: str, initial_catego
         "result_image": result_image,
         "result_image": result_image,
         "result_text": result_text,
+        "sample_gallery": sample_gallery,
     }