gradio_combine.py

import torch
import numpy as np
from PIL import Image
import cv2
import imutils
import os
import sys
import time
import colorsys
from scipy import ndimage
import gradio as gr
from huggingface_hub import hf_hub_download

# Detectron2 imports
from detectron2.config import get_cfg
from detectron2.projects.deeplab import add_deeplab_config
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor as DetectronPredictor
from detectron2 import model_zoo

# OneFormer imports
from oneformer import (
    add_oneformer_config,
    add_common_config,
    add_swin_config,
    add_dinat_config,
)
from demo.defaults import DefaultPredictor as OneFormerPredictor
from demo.visualizer import Visualizer, ColorMode

# NeuroNest contrast detection imports
from utils.contrast_detector import ContrastDetector
from utils.luminance_contrast import LuminanceContrastDetector
from utils.hue_contrast import HueContrastDetector
from utils.saturation_contrast import SaturationContrastDetector
from utils.combined_contrast import CombinedContrastDetector

# Set threads for CPU optimization
torch.set_num_threads(4)

########################################
# GLOBAL CONFIGURATIONS
########################################

# OneFormer configurations
KEY_DICT = {"ADE20K (150 classes)": "ade20k"}

SWIN_CFG_DICT = {
    "ade20k": "configs/ade20k/oneformer_swin_large_IN21k_384_bs16_160k.yaml",
}

SWIN_MODEL_DICT = {
    "ade20k": hf_hub_download(
        repo_id="shi-labs/oneformer_ade20k_swin_large",
        filename="250_16_swin_l_oneformer_ade20k_160k.pth"
    )
}

DINAT_CFG_DICT = {
    "ade20k": "configs/ade20k/oneformer_dinat_large_IN21k_384_bs16_160k.yaml",
}

DINAT_MODEL_DICT = {
    "ade20k": hf_hub_download(
        repo_id="shi-labs/oneformer_ade20k_dinat_large",
        filename="250_16_dinat_l_oneformer_ade20k_160k.pth"
    )
}

MODEL_DICT = {"DiNAT-L": DINAT_MODEL_DICT, "Swin-L": SWIN_MODEL_DICT}
CFG_DICT = {"DiNAT-L": DINAT_CFG_DICT, "Swin-L": SWIN_CFG_DICT}
WIDTH_DICT = {"ade20k": 640}

# Contrast detector mapping
CONTRAST_DETECTORS = {
    "Luminance (WCAG)": LuminanceContrastDetector(),
    "Hue": HueContrastDetector(),
    "Saturation": SaturationContrastDetector(),
    "Combined": CombinedContrastDetector()
}

# Device configuration
cpu_device = torch.device("cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Model storage
ONEFORMER_PREDICTORS = {
    "DiNAT-L": {"ADE20K (150 classes)": None},
    "Swin-L": {"ADE20K (150 classes)": None}
}

ONEFORMER_METADATA = {
    "DiNAT-L": {"ADE20K (150 classes)": None},
    "Swin-L": {"ADE20K (150 classes)": None}
}

########################################
# MASK R-CNN SETUP AND FUNCTIONS
########################################

def load_maskrcnn_model(weights_path, device="cuda", threshold=0.5):
    """Load Mask R-CNN model for blackspot detection"""
    cfg = get_cfg()
    cfg.merge_from_file(
        model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
    )
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # [Floors, blackspot]
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
    cfg.MODEL.WEIGHTS = weights_path
    # Fix: Convert torch.device to string
    cfg.MODEL.DEVICE = str(device) if isinstance(device, torch.device) else device
    return DetectronPredictor(cfg)
def postprocess_blackspot_masks(im, instances, show_floor=True, show_blackspot=True):
    """Extract floor and blackspot masks from Mask R-CNN predictions"""
    height, width = im.shape[:2]
    pred_classes = instances.pred_classes.cpu().numpy()
    pred_masks = instances.pred_masks.cpu().numpy()
    
    combined_floor_mask = np.zeros((height, width), dtype=bool)
    final_blackspot = np.zeros((height, width), dtype=bool)
    
    for cls_id, mask in zip(pred_classes, pred_masks):
        if cls_id == 0 and show_floor:  # Floor class
            combined_floor_mask |= mask
        elif cls_id == 1 and show_blackspot:  # Blackspot class
            final_blackspot |= mask
    
    return combined_floor_mask.astype(np.uint8), final_blackspot.astype(np.uint8)

########################################
# ONEFORMER SETUP AND FUNCTIONS
########################################

def setup_oneformer_modules():
    """Initialize OneFormer models"""
    for dataset in ["ADE20K (150 classes)"]:
        for backbone in ["DiNAT-L", "Swin-L"]:
            cfg = setup_oneformer_cfg(dataset, backbone)
            metadata = MetadataCatalog.get(
                cfg.DATASETS.TEST_PANOPTIC[0] if len(cfg.DATASETS.TEST_PANOPTIC) else "__unused"
            )
            ONEFORMER_PREDICTORS[backbone][dataset] = OneFormerPredictor(cfg)
            ONEFORMER_METADATA[backbone][dataset] = metadata

def setup_oneformer_cfg(dataset, backbone):
    """Setup OneFormer configuration"""
    cfg = get_cfg()
    add_deeplab_config(cfg)
    add_common_config(cfg)
    add_swin_config(cfg)
    add_oneformer_config(cfg)
    add_dinat_config(cfg)
    dataset = KEY_DICT[dataset]
    cfg_path = CFG_DICT[backbone][dataset]
    cfg.merge_from_file(cfg_path)
    cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    cfg.MODEL.WEIGHTS = MODEL_DICT[backbone][dataset]
    cfg.freeze()
    return cfg

def semantic_run(img, predictor, metadata):
    """Run OneFormer semantic segmentation"""
    visualizer = Visualizer(img[:, :, ::-1], metadata=metadata, instance_mode=ColorMode.IMAGE)
    predictions = predictor(img, "semantic")
    out = visualizer.draw_sem_seg(
        predictions["sem_seg"].argmax(dim=0).to(cpu_device), alpha=0.5
    )
    return out, predictions["sem_seg"].argmax(dim=0).to(cpu_device).numpy()

########################################
# INTEGRATED ANALYSIS FUNCTION
########################################

def integrated_analysis(image_path, 
                       # Blackspot detection parameters
                       blackspot_threshold, show_floor, show_blackspot,
                       # Contrast detection parameters
                       enable_contrast, backbone, contrast_method, contrast_threshold):
    """
    Perform integrated analysis with both blackspot detection and contrast analysis
    """
    # Read the image
    im = cv2.imread(image_path)
    if im is None:
        return "Error: could not read image!", None, None, None
    
    # Resize for OneFormer if contrast analysis is enabled
    if enable_contrast:
        width = WIDTH_DICT["ade20k"]
        im_resized = imutils.resize(im, width=width)
    else:
        im_resized = im
    
    # Part 1: Blackspot Detection with Mask R-CNN
    blackspot_text = []
    blackspot_viz = None
    
    if show_floor or show_blackspot:
        weights_path = "./output_floor_blackspot/model_0004999.pth"
        maskrcnn_predictor = load_maskrcnn_model(weights_path, device, blackspot_threshold)
        
        # Run blackspot detection
        outputs = maskrcnn_predictor(im)
        instances = outputs["instances"]
        
        # Post-process masks
        floor_mask, blackspot_mask = postprocess_blackspot_masks(im, instances, show_floor, show_blackspot)
        
        # Create visualization
        blackspot_overlay = im.copy()
        overlay = np.zeros_like(im)
        
        if show_floor:
            overlay[floor_mask > 0] = (0, 255, 0)  # Green for floor
        if show_blackspot:
            overlay[blackspot_mask > 0] = (0, 0, 255)  # Red for blackspot
        
        blackspot_overlay = cv2.addWeighted(im, 1.0, overlay, 0.5, 0)
        blackspot_viz = Image.fromarray(cv2.cvtColor(blackspot_overlay, cv2.COLOR_BGR2RGB))
        
        # Calculate statistics
        blackspot_area = int(blackspot_mask.sum())
        floor_area = int(floor_mask.sum())
        
        blackspot_text.append(f"### Blackspot Detection Results")
        blackspot_text.append(f"**Threshold:** {blackspot_threshold:.2f}")
        
        if show_floor:
            blackspot_text.append(f"**Floor area:** {floor_area} pixels")
        if show_blackspot:
            blackspot_text.append(f"**Blackspot area:** {blackspot_area} pixels")
            if floor_area > 0 and show_floor:
                percentage = (blackspot_area / floor_area) * 100
                blackspot_text.append(f"**Blackspot coverage:** {percentage:.2f}% of floor area")
    
    # Part 2: Contrast Analysis with OneFormer
    segmentation_viz = None
    contrast_viz = None
    contrast_text = []
    
    if enable_contrast:
        dataset = "ADE20K (150 classes)"
        predictor = ONEFORMER_PREDICTORS[backbone][dataset]
        metadata = ONEFORMER_METADATA[backbone][dataset]
        
        # Get segmentation
        out, seg_mask = semantic_run(im_resized, predictor, metadata)
        segmentation_viz = Image.fromarray(out.get_image())
        
        # Analyze contrast
        img_rgb = cv2.cvtColor(im_resized, cv2.COLOR_BGR2RGB)
        detector = CONTRAST_DETECTORS[contrast_method]
        contrast_image, problem_areas, stats = detector.analyze(
            img_rgb, seg_mask, contrast_threshold
        )
        
        contrast_viz = Image.fromarray(contrast_image)
        
        # Create stats text
        contrast_text.append(f"### Contrast Analysis Results")
        contrast_text.append(f"**Method:** {contrast_method}")
        contrast_text.append(f"**Threshold:** {contrast_threshold:.2f}")
        contrast_text.append(f"**Problem Areas:** {stats['problem_count']}")
        
        if 'min_contrast' in stats:
            contrast_text.append(f"**Min Contrast:** {stats['min_contrast']:.2f}")
        if 'max_contrast' in stats:
            contrast_text.append(f"**Max Contrast:** {stats['max_contrast']:.2f}")
        if 'average_contrast' in stats:
            contrast_text.append(f"**Average Contrast:** {stats['average_contrast']:.2f}")
    
    # Combine results
    combined_text = []
    if blackspot_text:
        combined_text.extend(blackspot_text)
    if contrast_text:
        if blackspot_text:
            combined_text.append("\n")
        combined_text.extend(contrast_text)
    
    return "\n".join(combined_text), blackspot_viz, segmentation_viz, contrast_viz

########################################
# GRADIO INTERFACE
########################################

# Initialize models
print("Initializing OneFormer models...")
setup_oneformer_modules()

title = "NeuroNest: Integrated Blackspot & Contrast Detection"
description = """
This integrated system combines:
1. **Blackspot Detection**: Uses Mask R-CNN to detect blackspots on floors
2. **Contrast Analysis**: Uses OneFormer segmentation to analyze contrast between objects

Both analyses help identify potential accessibility issues for individuals with Alzheimer's disease.
"""

# Create the Gradio interface
with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input image
            image_input = gr.Image(label="Input Image", type="filepath")
            
            # Blackspot detection controls
            with gr.Accordion("Blackspot Detection Settings", open=True):
                blackspot_threshold = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.5, step=0.05,
                    label="Blackspot Detection Threshold"
                )
                with gr.Row():
                    show_floor = gr.Checkbox(value=True, label="Show Floor")
                    show_blackspot = gr.Checkbox(value=True, label="Show Blackspots")
            
            # Contrast analysis controls
            with gr.Accordion("Contrast Analysis Settings", open=True):
                enable_contrast = gr.Checkbox(value=True, label="Enable Contrast Analysis")
                backbone = gr.Radio(
                    choices=["Swin-L", "DiNAT-L"], 
                    value="Swin-L", 
                    label="OneFormer Backbone"
                )
                contrast_method = gr.Radio(
                    choices=["Luminance (WCAG)", "Hue", "Saturation", "Combined"],
                    value="Luminance (WCAG)",
                    label="Contrast Detection Method"
                )
                contrast_threshold = gr.Slider(
                    minimum=1.0, maximum=10.0, value=4.5, step=0.1,
                    label="Contrast Threshold"
                )
            
            analyze_btn = gr.Button("Analyze", variant="primary")
        
        with gr.Column(scale=2):
            # Output displays
            with gr.Tabs():
                with gr.Tab("Analysis Report"):
                    analysis_text = gr.Textbox(label="Analysis Results", lines=10)
                
                with gr.Tab("Blackspot Detection"):
                    blackspot_output = gr.Image(label="Blackspot Visualization")
                
                with gr.Tab("Segmentation"):
                    segmentation_output = gr.Image(label="OneFormer Segmentation")
                
                with gr.Tab("Contrast Analysis"):
                    contrast_output = gr.Image(label="Contrast Visualization")
    
    # Connect the interface
    analyze_btn.click(
        fn=integrated_analysis,
        inputs=[
            image_input,
            blackspot_threshold, show_floor, show_blackspot,
            enable_contrast, backbone, contrast_method, contrast_threshold
        ],
        outputs=[
            analysis_text, blackspot_output, segmentation_output, contrast_output
        ]
    )
    
    # Examples
    gr.Examples(
        examples=[
            ["examples/indoor_room.jpg", 0.5, True, True, True, "Swin-L", "Luminance (WCAG)", 4.5],
            ["examples/living_room.jpg", 0.7, True, True, True, "DiNAT-L", "Combined", 3.0],
        ],
        inputs=[
            image_input,
            blackspot_threshold, show_floor, show_blackspot,
            enable_contrast, backbone, contrast_method, contrast_threshold
        ]
    )

if __name__ == "__main__":
    print(f"Launching integrated NeuroNest app on device: {device}")
    demo.queue().launch(server_name="0.0.0.0", share=True)