""" 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("""

Illumination-Invariant Identification of Wild and Farmed Larimichthys crocea

via Cross-Modal 3D Shape Priors

""") 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)