Updated UI to resize according to model preferred dimensions (640x640

gradio_test.py

@@ -52,35 +52,54 @@ ONEFORMER_CONFIG = {
         "swin_cfg": "configs/ade20k/oneformer_swin_large_IN21k_384_bs16_160k.yaml",
         "swin_model": "shi-labs/oneformer_ade20k_swin_large",
         "swin_file": "250_16_swin_l_oneformer_ade20k_160k.pth",
-        "width": 640
+        "process_size": 640,
+        "max_size": 2560
     }
 }
 
 BLACKSPOT_MODEL_REPO = "sww35/neuronest-blackspot"
 BLACKSPOT_MODEL_FILE = "model_0004999.pth"
 
-MAX_DISPLAY_SIZE = 1280
-MIN_DISPLAY_SIZE = 800
+DISPLAY_MAX_WIDTH = 1920
+DISPLAY_MAX_HEIGHT = 1080
 
 from universal_contrast_analyzer import UniversalContrastAnalyzer
 
-def resize_for_display(image: np.ndarray, max_size: int = MAX_DISPLAY_SIZE, min_size: int = MIN_DISPLAY_SIZE) -> np.ndarray:
+def resize_image_for_processing(image: np.ndarray, target_size: int = 640, max_size: int = 2560) -> Tuple[np.ndarray, float]:
     h, w = image.shape[:2]
-    if max(h, w) < min_size:
-        scale = min_size / max(h, w)
-    elif max(h, w) > max_size:
+    scale = target_size / min(h, w)
+    if scale * max(h, w) > max_size:
         scale = max_size / max(h, w)
-    else:
-        return image
     new_w = int(w * scale)
     new_h = int(h * scale)
-    return cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    new_w = (new_w // 32) * 32
+    new_h = (new_h // 32) * 32
+    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
+    return resized, scale
+
+def resize_mask_to_original(mask: np.ndarray, original_size: Tuple[int, int]) -> np.ndarray:
+    return cv2.resize(mask.astype(np.uint8), (original_size[1], original_size[0]), interpolation=cv2.INTER_NEAREST)
+
+def prepare_display_image(image: np.ndarray, max_width: int = DISPLAY_MAX_WIDTH, max_height: int = DISPLAY_MAX_HEIGHT) -> np.ndarray:
+    h, w = image.shape[:2]
+    scale = 1.0
+    if w > max_width:
+        scale = max_width / w
+    if h * scale > max_height:
+        scale = max_height / h
+    if scale < 1.0:
+        new_w = int(w * scale)
+        new_h = int(h * scale)
+        return cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
+    return image
 
 class OneFormerManager:
     def __init__(self):
         self.predictor = None
         self.metadata = None
         self.initialized = False
+        self.process_size = ONEFORMER_CONFIG["ADE20K"]["process_size"]
+        self.max_size = ONEFORMER_CONFIG["ADE20K"]["max_size"]
 
     def initialize(self, backbone: str = "swin"):
         if not ONEFORMER_AVAILABLE:
@@ -113,31 +132,25 @@ class OneFormerManager:
             logger.error(f"Failed to initialize OneFormer: {e}")
             return False
 
-    def semantic_segmentation(self, image: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def semantic_segmentation(self, image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
         if not self.initialized:
             raise RuntimeError("OneFormer not initialized")
-        original_h, original_w = image.shape[:2]
-        width = ONEFORMER_CONFIG["ADE20K"]["width"]
-        scale = width / original_w
-        new_h = int(original_h * scale)
-        image_resized = cv2.resize(image, (width, new_h), interpolation=cv2.INTER_LINEAR)
-        predictions = self.predictor(image_resized, "semantic")
-        seg_mask = predictions["sem_seg"].argmax(dim=0).cpu().numpy()
-        seg_mask_original = cv2.resize(
-            seg_mask.astype(np.uint8),
-            (original_w, original_h),
-            interpolation=cv2.INTER_NEAREST
-        )
+        original_size = (image.shape[0], image.shape[1])
+        image_processed, scale = resize_image_for_processing(image, self.process_size, self.max_size)
+        logger.info(f"Processing image at {image_processed.shape}, scale: {scale}")
+        predictions = self.predictor(image_processed, "semantic")
+        seg_mask_processed = predictions["sem_seg"].argmax(dim=0).cpu().numpy()
+        seg_mask_original = resize_mask_to_original(seg_mask_processed, original_size)
         visualizer = Visualizer(
             image[:, :, ::-1],
             metadata=self.metadata,
             instance_mode=ColorMode.IMAGE,
-            scale=1.2
+            scale=1.0
         )
-        vis_output = visualizer.draw_sem_seg(seg_mask_original, alpha=0.5)
+        vis_output = visualizer.draw_sem_seg(seg_mask_original, alpha=0.6)
         vis_image = vis_output.get_image()[:, :, ::-1]
-        vis_image_display = resize_for_display(vis_image)
-        return seg_mask_original, vis_image_display, image
+        vis_image_display = prepare_display_image(vis_image)
+        return seg_mask_original, vis_image_display
 
     def extract_floor_areas(self, segmentation: np.ndarray) -> np.ndarray:
         floor_mask = np.zeros_like(segmentation, dtype=bool)
@@ -275,7 +288,7 @@ class ImprovedBlackspotDetector:
         for mask in filtered_blackspot_masks:
             combined_blackspot |= mask
         visualization = self.create_visualization(image, floor_mask, combined_blackspot)
-        visualization_display = resize_for_display(visualization)
+        visualization_display = prepare_display_image(visualization)
         floor_area = int(np.sum(floor_mask))
         blackspot_area = int(np.sum(combined_blackspot))
         coverage_percentage = (blackspot_area / floor_area * 100) if floor_area > 0 else 0
@@ -304,12 +317,12 @@ class ImprovedBlackspotDetector:
         blackspot_contours, _ = cv2.findContours(
             blackspot_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
         )
-        cv2.drawContours(vis, blackspot_contours, -1, (255, 255, 0), 3)
+        cv2.drawContours(vis, blackspot_contours, -1, (255, 255, 0), 4)
         return vis
 
     def _empty_results(self, image: np.ndarray) -> Dict:
         empty_mask = np.zeros(image.shape[:2], dtype=bool)
-        visualization_display = resize_for_display(image)
+        visualization_display = prepare_display_image(image)
         return {
             'visualization': visualization_display,
             'floor_mask': empty_mask,
@@ -351,7 +364,7 @@ class NeuroNestApp:
         if not self.initialized:
             return {"error": "Application not properly initialized"}
         try:
-            image = cv2.imread(image_path)
+            image = cv2.imread(image_path, cv2.IMREAD_COLOR)
             if image is None:
                 return {"error": "Could not load image"}
             image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -364,7 +377,7 @@ class NeuroNestApp:
                 'statistics': {}
             }
             logger.info("Running semantic segmentation...")
-            seg_mask, seg_visualization, original_sized_image = self.oneformer.semantic_segmentation(image_rgb)
+            seg_mask, seg_visualization = self.oneformer.semantic_segmentation(image_rgb)
             results['segmentation'] = {
                 'visualization': seg_visualization,
                 'mask': seg_mask
@@ -387,7 +400,7 @@ class NeuroNestApp:
                     contrast_results = self.contrast_analyzer.analyze_contrast(
                         image_rgb, seg_mask
                     )
-                    contrast_viz_display = resize_for_display(contrast_results['visualization'])
+                    contrast_viz_display = prepare_display_image(contrast_results['visualization'])
                     contrast_results['visualization'] = contrast_viz_display
                     results['contrast'] = contrast_results
                     logger.info("Contrast analysis completed")
@@ -535,105 +548,124 @@ def create_gradio_interface():
     - **Floor-Only Blackspot Detection**: Locates dangerous dark areas on walking surfaces
     - **Universal Contrast Analysis**: Evaluates visibility between ALL adjacent objects
 
-    *Following WCAG 2.1 guidelines for visual accessibility*
+    *Following WCAG 2.1 guidelines for visual accessibility  | Upload a Picture. Click 'Analyze Environment'.Then scroll down.*
     """
     with gr.Blocks(css="""
-        .image-container img {
-            max-width: 100%;
-            height: auto;
-            object-fit: contain;
+        .container { max-width: 100%; margin: auto; padding: 20px; }
+        .image-output { margin: 20px 0; }
+        .image-output img { 
+            width: 100%; 
+            height: auto; 
+            max-width: 1920px; 
+            margin: 0 auto; 
+            display: block;
+            border: 1px solid #ddd;
+            border-radius: 8px;
         }
-    """) as interface:
-        gr.Markdown(f"# {title}")
-        gr.Markdown(description)
-        if not blackspot_ok:
-            gr.Markdown("""
-            ⚠️ **Note:** Blackspot detection model not available. 
-            To enable blackspot detection, upload the model to HuggingFace or ensure it's in the local directory.
-            """)
-        with gr.Row():
-            enable_blackspot = gr.Checkbox(
-                value=blackspot_ok,
-                label="Enable Floor Blackspot Detection",
-                interactive=blackspot_ok
-            )
-            blackspot_threshold = gr.Slider(
-                minimum=0.1,
-                maximum=0.9,
-                value=0.5,
-                step=0.05,
-                label="Blackspot Sensitivity",
-                visible=blackspot_ok
-            )
-            enable_contrast = gr.Checkbox(
-                value=True,
-                label="Enable Universal Contrast Analysis"
-            )
-            contrast_threshold = gr.Slider(
-                minimum=3.0,
-                maximum=7.0,
-                value=4.5,
-                step=0.1,
-                label="WCAG Contrast Threshold"
-            )
-        with gr.Row():
-            with gr.Column(scale=1):
-                image_input = gr.Image(
-                    label="📸 Upload Room Image",
-                    type="filepath",
-                    height=400
-                )
-            with gr.Column(scale=1):
-                analyze_button = gr.Button(
-                    "🔍 Analyze Environment",
-                    variant="primary",
-                    size="lg"
-                )
-        with gr.Row(elem_classes="image-container"):
+        .controls-row { margin-bottom: 30px; background: #f5f5f5; padding: 20px; border-radius: 8px; }
+        .main-button { height: 80px !important; font-size: 1.3em !important; font-weight: bold !important; }
+        .report-box { max-width: 1200px; margin: 30px auto; padding: 30px; background: #f9f9f9; border-radius: 8px; }
+        h2 { margin-top: 40px; margin-bottom: 20px; color: #333; }
+    """, theme=gr.themes.Base()) as interface:
+        with gr.Column(elem_classes="container"):
+            gr.Markdown(f"# {title}")
+            gr.Markdown(description)
+            if not blackspot_ok:
+                gr.Markdown("""
+                ⚠️ **Note:** Blackspot detection model not available. 
+                To enable blackspot detection, upload the model to HuggingFace or ensure it's in the local directory.
+                """)
+            with gr.Row(elem_classes="controls-row"):
+                with gr.Column(scale=1):
+                    enable_blackspot = gr.Checkbox(
+                        value=blackspot_ok,
+                        label="Enable Floor Blackspot Detection",
+                        interactive=blackspot_ok
+                    )
+                    blackspot_threshold = gr.Slider(
+                        minimum=0.1,
+                        maximum=0.9,
+                        value=0.5,
+                        step=0.05,
+                        label="Blackspot Sensitivity",
+                        visible=blackspot_ok
+                    )
+                with gr.Column(scale=1):
+                    enable_contrast = gr.Checkbox(
+                        value=True,
+                        label="Enable Universal Contrast Analysis"
+                    )
+                    contrast_threshold = gr.Slider(
+                        minimum=3.0,
+                        maximum=7.0,
+                        value=4.5,
+                        step=0.1,
+                        label="WCAG Contrast Threshold"
+                    )
+            with gr.Row():
+                with gr.Column(scale=2):
+                    image_input = gr.Image(
+                        label="📸 Upload Room Image",
+                        type="filepath",
+                        height=500
+                    )
+                with gr.Column(scale=1):
+                    analyze_button = gr.Button(
+                        "🔍 Analyze Environment",
+                        variant="primary",
+                        elem_classes="main-button"
+                    )
+            gr.Markdown("---")
+            gr.Markdown("## 🎯 Segmented Objects")
             seg_display = gr.Image(
-                label="🎯 Segmented Objects",
+                label=None,
                 interactive=False,
-                show_label=True,
-                height=600
-            )
-            blackspot_display = gr.Image(
-                label="⚫ Blackspot Detection",
-                interactive=False,
-                visible=blackspot_ok,
-                show_label=True,
-                height=600
+                show_label=False,
+                elem_classes="image-output"
             )
+            if blackspot_ok:
+                gr.Markdown("## ⚫ Blackspot Detection")
+                blackspot_display = gr.Image(
+                    label=None,
+                    interactive=False,
+                    show_label=False,
+                    elem_classes="image-output"
+                )
+            else:
+                blackspot_display = gr.Image(visible=False)
+            gr.Markdown("## 🎨 Contrast Analysis")
             contrast_display = gr.Image(
-                label="🎨 Contrast Analysis",
+                label=None,
                 interactive=False,
-                show_label=True,
-                height=600
+                show_label=False,
+                elem_classes="image-output"
             )
-        analysis_report = gr.Markdown(
-            value="Upload an image and click 'Analyze Environment' to begin.",
-            elem_classes="report-container"
-        )
-        analyze_button.click(
-            fn=analyze_wrapper,
-            inputs=[
-                image_input,
-                blackspot_threshold,
-                contrast_threshold,
-                enable_blackspot,
-                enable_contrast
-            ],
-            outputs=[
-                seg_display,
-                blackspot_display,
-                contrast_display,
-                analysis_report
-            ]
-        )
-        gr.Markdown("""
-            ---
-            **NeuroNest** v2.0 - Enhanced with floor-only blackspot detection and universal contrast analysis  
-            *Creating safer environments for cognitive health through AI*
-            """)
+            gr.Markdown("---")
+            analysis_report = gr.Markdown(
+                value="Upload an image and click 'Analyze Environment' to begin.",
+                elem_classes="report-box"
+            )
+            analyze_button.click(
+                fn=analyze_wrapper,
+                inputs=[
+                    image_input,
+                    blackspot_threshold,
+                    contrast_threshold,
+                    enable_blackspot,
+                    enable_contrast
+                ],
+                outputs=[
+                    seg_display,
+                    blackspot_display,
+                    contrast_display,
+                    analysis_report
+                ]
+            )
+            gr.Markdown("""
+                ---
+                **NeuroNest** v2.0 - Enhanced with floor-only blackspot detection and universal contrast analysis  
+                *Creating safer environments for cognitive health through AI*
+                """)
     return interface
 
 if __name__ == "__main__":