Blame for app.py

"""
Hugging Face Spaces App for SGD-Net Fish Classification
Enhanced with Grad-CAM visualization
"""
import gradio as gr
import torch
import torch.nn.functional as F
import numpy as np
import cv2
from pathlib import Path
from ultralytics import YOLO
import albumentations as A
from albumentations.pytorch import ToTensorV2
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

from models.sgd_net import ResNetSGDNet
from gradcam import GradCAM, get_target_layer

# Model paths
YOLO_MODEL_PATH = "yolo_model.pt"
SGD_NET_CHECKPOINT = "sgd_net_checkpoint.pth"

# SGD-Net configuration
SGD_NET_CONFIG = {
    "backbone": "resnet50",
    "num_classes": 2,
    "pretrained_path": None,
    "use_shape_attention": False,
    "use_attention_gate": True
}

# Global variables for models
yolo_model = None
sgd_net_model = None
gradcam = None
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Preprocessing transform
transform = A.Compose([
    A.Resize(256, 256),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2()
])


def load_models():
    """Load YOLO and SGD-Net models"""
    global yolo_model, sgd_net_model, gradcam
    
    # Load YOLO model
    print("Loading YOLO model...")
    try:
        yolo_model = YOLO(YOLO_MODEL_PATH)
        print("YOLO model loaded successfully")
    except Exception as e:
        print(f"Error loading YOLO model: {e}")
        raise
    
    # Load SGD-Net model
    print("Loading SGD-Net model...")
    sgd_net_model = ResNetSGDNet(
        backbone_name=SGD_NET_CONFIG["backbone"],
        num_classes=SGD_NET_CONFIG["num_classes"],
        pretrained=True,
        pretrained_path=SGD_NET_CONFIG.get("pretrained_path"),
        use_shape_attention=SGD_NET_CONFIG["use_shape_attention"],
        use_attention_gate=SGD_NET_CONFIG["use_attention_gate"]
    )
    
    # Load checkpoint
    print(f"Loading checkpoint from: {SGD_NET_CHECKPOINT}")
    checkpoint = torch.load(SGD_NET_CHECKPOINT, map_location=device, weights_only=False)
    
    # Handle different checkpoint formats
    if isinstance(checkpoint, dict):
        if 'model_state_dict' in checkpoint:
            state_dict = checkpoint['model_state_dict']
        elif 'state_dict' in checkpoint:
            state_dict = checkpoint['state_dict']
        else:
            state_dict = checkpoint
    else:
        state_dict = checkpoint
    
    # Filter mismatched keys
    model_state_dict = sgd_net_model.state_dict()
    filtered_state_dict = {}
    for k, v in state_dict.items():
        if k in model_state_dict and model_state_dict[k].shape == v.shape:
            filtered_state_dict[k] = v
        elif k.startswith('module.'):
            k_new = k[7:]
            if k_new in model_state_dict and model_state_dict[k_new].shape == v.shape:
                filtered_state_dict[k_new] = v
    
    sgd_net_model.load_state_dict(filtered_state_dict, strict=False)
    sgd_net_model = sgd_net_model.to(device)
    sgd_net_model.eval()
    print("SGD-Net model loaded successfully")
    
    # Initialize Grad-CAM
    print("Initializing Grad-CAM...")
    target_layer = get_target_layer(sgd_net_model)
    if target_layer is not None:
        gradcam = GradCAM(sgd_net_model, target_layer)
        print("Grad-CAM initialized successfully")
    else:
        print("Warning: Could not find target layer for Grad-CAM")


def get_yolo_mask_and_confidence(rgb_image, conf_thres=0.5):
    """
    Generate mask using YOLO and return confidence
    
    Returns:
        mask: Binary mask (H, W) uint8, 255=foreground, 0=background
        confidence: Maximum confidence, returns 0.0 if no detection
    """
    results = yolo_model.predict(rgb_image, conf=conf_thres, verbose=False)
    result = results[0]
    
    confidence = 0.0
    mask = np.zeros(rgb_image.shape[:2], dtype=np.uint8)
    
    if result.masks is not None and len(result.masks) > 0:
        boxes = result.boxes
        if len(boxes) > 0:
            confs = boxes.conf.cpu().numpy()
            best_idx = int(np.argmax(confs))
            confidence = float(confs[best_idx])
            
            mask_data = result.masks[best_idx].data.cpu().numpy()
            if mask_data.ndim == 3:
                mask_data = mask_data[0]
            
            mask = (mask_data * 255).astype(np.uint8)
            
            if mask.shape != rgb_image.shape[:2]:
                mask = cv2.resize(mask, (rgb_image.shape[1], rgb_image.shape[0]), 
                                interpolation=cv2.INTER_NEAREST)
    
    return mask, confidence


def crop_image_by_mask(rgb_image, mask, expand_ratio=0.1):
    """
    Crop image and mask based on mask bounding box
    
    Args:
        rgb_image: RGB image (H, W, 3) numpy array
        mask: Binary mask (H, W) numpy array, 255=foreground, 0=background
        expand_ratio: Expansion ratio for bounding box
    
    Returns:
        rgb_cropped: Cropped RGB image
        mask_cropped: Cropped mask
    """
    mask_binary = mask > 0
    if mask_binary.sum() > 0:
        coords = np.column_stack(np.where(mask_binary))
        y_min, x_min = coords.min(axis=0)
        y_max, x_max = coords.max(axis=0)
        
        H_img, W_img = rgb_image.shape[:2]
        margin_y = int((y_max - y_min) * expand_ratio)
        margin_x = int((x_max - x_min) * expand_ratio)
        
        y_min = max(0, y_min - margin_y)
        y_max = min(H_img, y_max + margin_y)
        x_min = max(0, x_min - margin_x)
        x_max = min(W_img, x_max + margin_x)
        
        # Crop RGB and mask based on bounding box
        rgb_cropped = rgb_image[y_min:y_max, x_min:x_max].copy()
        mask_cropped = mask[y_min:y_max, x_min:x_max].copy()
    else:
        # If mask is empty, use full image
        rgb_cropped = rgb_image
        mask_cropped = mask
    
    return rgb_cropped, mask_cropped


def preprocess_image(rgb_image, mask):
    """
    Preprocess image: crop by mask, resize and normalize
    
    Args:
        rgb_image: RGB image (H, W, 3) numpy array
        mask: Binary mask (H, W) numpy array
    
    Returns:
        tensor: Preprocessed image tensor (1, 3, 256, 256)
        rgb_cropped: Cropped RGB image (for visualization)
        mask_cropped: Cropped mask (for visualization)
    """
    # Step 1: Crop image and mask based on mask bounding box
    rgb_cropped, mask_cropped = crop_image_by_mask(rgb_image, mask, expand_ratio=0.1)
    
    # Step 2: Resize and normalize
    transformed = transform(image=rgb_cropped)
    rgb_tensor = transformed['image']
    rgb_tensor = rgb_tensor.unsqueeze(0)
    
    return rgb_tensor, rgb_cropped, mask_cropped


def generate_gradcam_heatmap(cam, original_image, mask):
    """
    Generate Grad-CAM heatmap overlaid on original image (only within mask)
    
    Args:
        cam: CAM heatmap (H, W) numpy array, values in [0, 1]
        original_image: Original image (H, W, 3) numpy array
        mask: Binary mask (H, W) numpy array, 255=foreground, 0=background
    
    Returns:
        overlay: Overlaid image (H, W, 3) numpy array
    """
    # Resize CAM to original image size
    h, w = original_image.shape[:2]
    cam_resized = cv2.resize(cam, (w, h), interpolation=cv2.INTER_LINEAR)
    
    # Resize mask to match image size if needed
    if mask.shape[:2] != (h, w):
        mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    else:
        mask_resized = mask
    
    # Normalize mask to [0, 1]
    mask_normalized = (mask_resized > 0).astype(np.float32)
    
    # Apply mask to CAM (only show attention within mask)
    cam_masked = cam_resized * mask_normalized
    
    # Create heatmap colormap (jet colormap)
    cam_colored = plt.cm.jet(cam_masked)[:, :, :3]  # (H, W, 3) RGB
    cam_colored = (cam_colored * 255).astype(np.uint8)
    
    # Overlay: original * (1 - alpha) + heatmap * alpha, but only within mask
    alpha = 0.4
    mask_3d = np.stack([mask_normalized] * 3, axis=2)  # (H, W, 3)
    
    overlay = (original_image.astype(np.float32) * (1 - alpha * mask_3d) + 
               cam_colored.astype(np.float32) * (alpha * mask_3d)).astype(np.uint8)
    
    return overlay


def classify_fish(image):
    """
    Classify fish image
    
    Args:
        image: PIL Image or numpy array
    
    Returns:
        tuple: (result_text, rgb_image, mask_image, gradcam_image)
    """
    if image is None:
        return "Please upload an image", None, None, None
    
    # Convert to numpy array (RGB format)
    if hasattr(image, 'convert'):
        rgb_image = np.array(image.convert('RGB'))
    else:
        rgb_image = np.array(image)
        if len(rgb_image.shape) == 3 and rgb_image.shape[2] == 4:  # RGBA
            rgb_image = rgb_image[:, :, :3]
    
    # Step 1: YOLO detection
    mask, confidence = get_yolo_mask_and_confidence(rgb_image, conf_thres=0.5)
    
    # Step 2: Check confidence
    if confidence < 0.5:
        result_text = f"**Detection Confidence: {confidence:.3f} < 0.5**\n\n**Result: Not Larimichthys crocea**"
        mask_display = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
        return result_text, rgb_image, mask_display, None
    
    # Step 3: Crop and preprocess image for classification
    rgb_tensor, rgb_cropped, mask_cropped = preprocess_image(rgb_image, mask)
    rgb_tensor = rgb_tensor.to(device)
    
    # Inference
    with torch.no_grad():
        cls_logits, _ = sgd_net_model(rgb_tensor)
        probs = F.softmax(cls_logits, dim=1)
        pred_class = cls_logits.argmax(dim=1).item()
        prob_wild = probs[0, 0].item()
        prob_farmed = probs[0, 1].item()
    
    # Generate Grad-CAM heatmap (only within mask)
    gradcam_image = None
    if gradcam is not None:
        try:
            cam = gradcam.generate_cam(rgb_tensor, class_idx=pred_class)
            # Use cropped image and mask for Grad-CAM visualization
            gradcam_image = generate_gradcam_heatmap(cam, rgb_cropped, mask_cropped)
        except Exception as e:
            print(f"Error generating Grad-CAM: {e}")
            gradcam_image = None
    
    # Generate result text
    class_name = "Wild" if pred_class == 0 else "Farmed"
    
    result_text = f"""**Detection Confidence: {confidence:.3f} ≥ 0.5**

**Classification Result: {class_name}**

**Probability Distribution:**
- Wild: **{prob_wild:.1%}**
- Farmed: **{prob_farmed:.1%}**"""
    
    # Convert mask to 3-channel for display (use cropped mask)
    mask_display = cv2.cvtColor(mask_cropped, cv2.COLOR_GRAY2RGB)
    
    # Use cropped image for display
    return result_text, rgb_cropped, mask_display, gradcam_image


# Create Gradio interface
def create_interface():
    # Custom CSS styles
    custom_css = """
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    .main-header {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 20px;
        border-radius: 10px;
        color: white;
        margin-bottom: 20px;
    }
    .result-box {
        background: #f8f9fa;
        border-radius: 8px;
        padding: 15px;
        border-left: 4px solid #667eea;
    }
    """
    
    with gr.Blocks(title="Illumination-Invariant Fish Classification", css=custom_css, theme=gr.themes.Soft()) as demo:
        # Title
        gr.Markdown("""
        <div class="main-header">
        <h1>Illumination-Invariant Identification of Wild and Farmed Larimichthys crocea</h1>
        <p style="font-size: 14px; margin-top: 10px;">via Cross-Modal 3D Shape Priors</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Upload Image",
                    type="numpy",
                    height=300
                )
                classify_btn = gr.Button(
                    "Classify",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=2):
                result_output = gr.Textbox(
                    label="Classification Result",
                    lines=8,
                    elem_classes=["result-box"]
                )
        
        with gr.Row():
            with gr.Column():
                image_output = gr.Image(
                    label="Original Image",
                    height=300
                )
            
            with gr.Column():
                mask_output = gr.Image(
                    label="Detection Mask",
                    height=300
                )
            
            with gr.Column():
                gradcam_output = gr.Image(
                    label="Grad-CAM Heatmap",
                    height=300
                )
        
        classify_btn.click(
            fn=classify_fish,
            inputs=image_input,
            outputs=[result_output, image_output, mask_output, gradcam_output]
        )
    
    return demo


if __name__ == "__main__":
    # Load models
    load_models()
    
    # Create and launch application
    demo = create_interface()
    demo.launch(share=False)