Spaces:
Sleeping
Sleeping
| """ | |
| 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) | |