tt

app.py

@@ -0,0 +1,575 @@
+# -*- coding: utf-8 -*-
+import os
+import base64
+import io
+import logging # For better logging
+# Import specific handlers and formatter
+from logging.handlers import RotatingFileHandler
+import traceback # For detailed exception logging
+from flask import Flask, request, jsonify, send_from_directory
+from flask_cors import CORS # To handle Cross-Origin requests from your frontend
+import torch
+import cv2
+import numpy as np
+import yaml
+from torchvision import transforms
+from transformers import SegformerForSemanticSegmentation
+from omegaconf import OmegaConf # Import OmegaConf itself
+import torch.nn.functional as F
+from werkzeug.utils import secure_filename # For safer filenames
+
+# --- Configuration ---
+# Use absolute paths for robustness
+BASE_DIR = os.path.dirname(os.path.abspath(__file__)) # Directory where this script is running
+
+# >>> Point this to your actual config file <<<
+CONFIG_PATH = os.path.join(BASE_DIR, "config/config.yaml") # Assuming config.yaml is in the same dir
+
+# >>> Point this to your actual checkpoint file <<<
+CHECKPOINT_PATH = "ckpt_000-vloss_0.4685_vf1_0.6469.ckpt"
+
+UPLOAD_FOLDER = os.path.join(BASE_DIR, 'uploads')
+RESULT_FOLDER = os.path.join(BASE_DIR, 'results')
+LOG_FILE_PATH = os.path.join(BASE_DIR, 'flask_app.log') # Define log file path
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp', 'tif', 'tiff'}
+
+# --- Logging Setup ---
+# Clear existing handlers from the root logger to avoid duplicates on reload
+logging.getLogger().handlers.clear()
+# Create formatter
+log_formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d - %(message)s')
+# Create Console Handler
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(log_formatter)
+# Create File Handler (using RotatingFileHandler for log rotation)
+file_handler = RotatingFileHandler(LOG_FILE_PATH, maxBytes=5*1024*1024, backupCount=3)
+file_handler.setFormatter(log_formatter)
+# Get the root logger and add handlers
+logger = logging.getLogger()
+logger.setLevel(logging.INFO) # Set minimum level for the logger (e.g., INFO, DEBUG)
+logger.addHandler(console_handler)
+logger.addHandler(file_handler)
+
+# --- Ensure upload and result directories exist ---
+try:
+    os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+    os.makedirs(RESULT_FOLDER, exist_ok=True)
+    logger.info(f"Ensured directories exist: {UPLOAD_FOLDER}, {RESULT_FOLDER}")
+except OSError as e:
+    logger.error(f"Error creating directories: {e}")
+    exit(1) # Exit if we can't create essential folders
+
+# --- Load Config ---
+config = None
+try:
+    # Load the YAML file using OmegaConf
+    config = OmegaConf.load(CONFIG_PATH)
+    # Note: We don't need OmegaConf.create() if loading directly from file
+    logger.info(f"Configuration loaded successfully from: {CONFIG_PATH}")
+    # Log some key values to confirm loading
+    logger.info(f"Config check: num_classes={config.data.num_classes}, model_name={config.training.model_name}")
+except FileNotFoundError:
+    logger.error(f"Configuration file not found: {CONFIG_PATH}")
+    exit(1)
+except Exception as e: # Catch broader errors during loading/parsing
+    logger.error(f"Error loading or parsing configuration file '{CONFIG_PATH}': {e}")
+    logger.error(traceback.format_exc())
+    exit(1)
+
+# --- Model Definition ---
+class InferenceModel(torch.nn.Module):
+    def __init__(self, model_config): # Use local name 'model_config'
+        super().__init__()
+        try:
+            # Access config values needed for model init
+            model_name = model_config.training.model_name
+            num_classes = model_config.data.num_classes
+            logger.info(f"Initializing SegformerForSemanticSegmentation with model='{model_name}', num_labels={num_classes}")
+            self.model = SegformerForSemanticSegmentation.from_pretrained(
+                model_name,
+                num_labels=num_classes,
+                ignore_mismatched_sizes=True # Important if fine-tuning head size differs
+            )
+            logger.info("Segformer model part initialized.")
+        except AttributeError as ae:
+             logger.error(f"Config error during model init: Missing key? {ae}")
+             logger.error(f"Check if 'training.model_name' and 'data.num_classes' exist in {CONFIG_PATH}")
+             raise # Re-raise error to stop execution
+        except Exception as e:
+            logger.error(f"Error initializing Segformer model from Hugging Face: {e}")
+            logger.error(traceback.format_exc())
+            raise # Re-raise error to stop execution
+
+    def forward(self, x):
+        # Expects pixel_values as input
+        outputs = self.model(pixel_values=x, return_dict=True)
+        # Upsample logits to original input size
+        logits = F.interpolate(
+            outputs.logits,
+            size=x.shape[-2:], # Get H, W from input tensor x
+            mode="bilinear",
+            align_corners=False
+        )
+        return logits
+
+# --- Utility Functions ---
+def num_to_rgb(num_arr, color_map_dict):
+    """Converts a label mask (numpy array) to an RGB color mask."""
+    single_layer = np.squeeze(num_arr)
+    output = np.zeros(num_arr.shape[:2] + (3,), dtype=np.uint8) # Initialize with uint8 zeros
+
+    # Expects color_map_dict to be a standard Python dict {int_label: [R, G, B]}
+    if not isinstance(color_map_dict, dict):
+         logger.error(f"Invalid color_map provided to num_to_rgb: {type(color_map_dict)}. Expected dict.")
+         return np.float32(output) / 255.0 # Return black float image
+
+    unique_labels = np.unique(single_layer)
+    for k in unique_labels:
+        label_key = int(k) # Ensure key is standard int for lookup
+        if label_key in color_map_dict:
+             # Assign color, ensure color value is appropriate (e.g., list/tuple of 3 ints)
+             color = color_map_dict[label_key]
+             if isinstance(color, (list, tuple)) and len(color) == 3:
+                 output[single_layer == k] = color
+             else:
+                 logger.warning(f"Invalid color format for label {label_key} in color map: {color}. Skipping.")
+        else:
+            if label_key != 0: # Often 0 is background, might not be in map
+                 logger.warning(f"Label Key {label_key} found in mask but not in provided color map.")
+            # Default color (e.g., black) is already set by np.zeros
+
+    return np.float32(output) / 255.0 # Return float32 RGB image [0, 1]
+
+def denormalize(tensor, mean, std):
+    """Denormalizes a torch tensor (CHW format)."""
+    # Expects standard Python lists/tuples for mean/std
+    if not isinstance(mean, (list, tuple)) or not isinstance(std, (list, tuple)):
+        logger.error(f"Mean ({type(mean)}) or std ({type(std)}) are not lists/tuples in denormalize.")
+        return None
+    # Input tensor expected shape: Batch, Channel, Height, Width (e.g., from dataloader or transform)
+    if tensor.dim() != 4: # B C H W
+        logger.error(f"Unexpected tensor dimension {tensor.dim()} in denormalize. Expected 4 (BCHW).")
+        # Attempt to add batch dim if it's 3D (CHW)
+        if tensor.dim() == 3:
+            logger.warning("Denormalize received 3D tensor, adding batch dimension.")
+            tensor = tensor.unsqueeze(0)
+        else:
+            return None # Cannot handle other dims
+
+    num_channels = tensor.shape[1] # Channel dimension
+    if len(mean) != num_channels or len(std) != num_channels:
+         logger.error(f"Mean/std length ({len(mean)}/{len(std)}) mismatch with tensor channels ({num_channels})")
+         return None
+
+    # Clone to avoid modifying original tensor
+    tensor = tensor.clone().cpu() # Work on CPU copy
+
+    # Denormalize each channel
+    for c in range(num_channels):
+        tensor[:, c, :, :] = tensor[:, c, :, :] * std[c] + mean[c] # Apply to all items in batch
+
+    # Clamp values, remove batch dimension, permute to HWC for display/saving
+    # Assumes we are processing one image at a time here for inference result
+    denormalized_img_tensor = torch.clamp(tensor.squeeze(0), 0, 1).permute(1, 2, 0)
+
+    return denormalized_img_tensor.numpy() # Convert to numpy array (HWC, float32, [0,1])
+
+# --- Load Model (Corrected Version) ---
+def load_trained_model(checkpoint_path, model_config):
+    """Loads the trained model from a checkpoint, handling potential key mismatches."""
+    try:
+        model_instance = InferenceModel(model_config) # Create model structure
+        logger.info(f"Attempting to load checkpoint from: {checkpoint_path}")
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"Checkpoint file not found at specified path: {checkpoint_path}")
+
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+        logger.info(f"Checkpoint loaded into memory. Type: {type(checkpoint)}")
+
+        # Extract the state dictionary - flexible based on common saving patterns
+        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            state_dict = checkpoint["state_dict"]
+            logger.info("Using 'state_dict' key from checkpoint.")
+        elif isinstance(checkpoint, dict):
+             # Assume the dict *is* the state_dict if 'state_dict' key is absent
+             state_dict = checkpoint
+             logger.info("Using checkpoint dictionary directly as state_dict (no 'state_dict' key found).")
+        else:
+            # Could be the model itself was saved directly (less common with frameworks)
+            logger.warning(f"Checkpoint is not a dictionary. Attempting to load directly into model (less common). Type was: {type(checkpoint)}")
+            # This path might need adjustment based on how the model was saved if not a state_dict
+            try:
+                 model_instance.load_state_dict(checkpoint) # Try loading directly
+                 logger.info("Loaded state_dict directly from checkpoint object.")
+                 model_instance.eval()
+                 return model_instance
+            except Exception as e:
+                 logger.error(f"Failed to load state_dict directly from checkpoint object: {e}")
+                 return None # Failed direct load
+
+        # --- Key Prefix Correction Logic ---
+        target_keys = set(model_instance.state_dict().keys())
+        loaded_keys = set(state_dict.keys())
+        if not loaded_keys: logger.warning("Loaded state_dict is empty!"); return None # Check if state_dict is empty
+        first_loaded_key = next(iter(loaded_keys), None)
+        first_target_key = next(iter(target_keys), None)
+        corrected_state_dict = {}
+        prefix_added = False
+
+        # Check if prefix 'model.' needs to be ADDED to loaded keys
+        if first_loaded_key and not first_loaded_key.startswith('model.') and \
+           first_target_key and first_target_key.startswith('model.'):
+            logger.warning("Checkpoint keys missing 'model.' prefix. Attempting to add it.")
+            prefix_added = True
+            keys_not_prefixed_properly = []
+            for k, v in state_dict.items():
+                new_key = f"model.{k}"
+                if new_key in target_keys: corrected_state_dict[new_key] = v
+                else: keys_not_prefixed_properly.append(k); corrected_state_dict[k] = v # Keep original if prefixed version not wanted
+            if keys_not_prefixed_properly: logger.warning(f"Keys kept without prefix (target doesn't expect): {keys_not_prefixed_properly}")
+            logger.info("Finished attempting prefix addition.")
+        # Check if prefix 'model.' needs to be REMOVED from loaded keys
+        elif first_loaded_key and first_loaded_key.startswith('model.') and \
+             first_target_key and not first_target_key.startswith('model.'):
+             logger.warning("Checkpoint keys HAVE 'model.' prefix, but target model DOES NOT. Attempting to remove it.")
+             prefix_added = False # Indicate we removed prefix, not added
+             keys_not_stripped_properly = []
+             for k, v in state_dict.items():
+                 if k.startswith('model.'):
+                      new_key = k.partition('model.')[2] # Get part after 'model.'
+                      if new_key in target_keys: corrected_state_dict[new_key] = v
+                      else: keys_not_stripped_properly.append(k); corrected_state_dict[k] = v # Keep original if stripped version not wanted
+                 else:
+                      # Keep keys that didn't have prefix anyway
+                      corrected_state_dict[k] = v
+             if keys_not_stripped_properly: logger.warning(f"Keys kept with prefix (target doesn't expect stripped): {keys_not_stripped_properly}")
+             logger.info("Finished attempting prefix removal.")
+        else:
+             logger.info("State dict keys seem to have correct prefix structure (or other mismatch). Using as is.")
+             corrected_state_dict = state_dict # Use the original dict
+
+        # --- Load the State Dictionary ---
+        logger.info("Attempting to load state_dict with strict=False for checking...")
+        missing_keys, unexpected_keys = model_instance.load_state_dict(corrected_state_dict, strict=False)
+
+        # Report detailed findings
+        final_msg = []
+        is_load_successful = True
+        if missing_keys:
+             final_msg.append(f"MISSING keys in checkpoint: {missing_keys}")
+             logger.error("CRITICAL FAILURE: Model is missing required keys.")
+             is_load_successful = False
+        if unexpected_keys:
+             final_msg.append(f"UNEXPECTED keys in checkpoint (exist in file but not in model): {unexpected_keys}")
+             # Decide if unexpected keys are acceptable
+             acceptable_unexpected = [k for k in unexpected_keys if k.endswith('num_batches_tracked')]
+             unacceptable_unexpected = [k for k in unexpected_keys if not k.endswith('num_batches_tracked')]
+             if unacceptable_unexpected:
+                 logger.error(f"CRITICAL FAILURE: Model received unacceptable unexpected keys: {unacceptable_unexpected}")
+                 is_load_successful = False
+             elif acceptable_unexpected:
+                 logger.warning(f"Ignoring acceptable unexpected keys: {acceptable_unexpected}")
+
+        if not is_load_successful:
+            logger.error(f"State dict loading failed. Issues: {'; '.join(final_msg)}")
+            return None # Failed to load properly
+
+        logger.info(f"State dictionary loaded successfully. Issues (if any): {final_msg if final_msg else 'None'}")
+        model_instance.eval() # Set to evaluation mode
+        logger.info(f"Model loading process complete for {checkpoint_path}")
+        return model_instance
+
+    except FileNotFoundError as fnf_error:
+        logger.error(f"{fnf_error}") # Log the specific FileNotFoundError message
+        return None
+    except Exception as e:
+        logger.error(f"Unexpected error during model loading: {e}")
+        logger.error(traceback.format_exc()) # Log full traceback
+        return None
+
+
+# --- Determine device & Load Model Globally ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info(f"Using device: {device}")
+# Load the model using the global config object
+model = load_trained_model(CHECKPOINT_PATH, config) # Pass the loaded config
+if model is None:
+    logger.critical("CRITICAL: Failed to load model. Application cannot continue.")
+    exit(1) # Critical error, stop the application
+model.to(device) # Move model to the appropriate device
+
+# --- Inference Pipeline (Corrected Config Handling) ---
+def run_inference_on_bytes(image_bytes, inference_model, model_config, device):
+    """Runs inference on image bytes, returns denormalized image, color mask, and overlay."""
+    try:
+        nparr = np.frombuffer(image_bytes, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        if img is None: logger.error("Failed cv2.imdecode."); return None, None, None
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        logger.debug("Image decoded and converted to RGB.")
+
+        # --- Preprocessing (with config conversion) ---
+        # Check necessary config attributes exist before conversion attempt
+        required_data_keys = ['image_size', 'mean', 'std', 'num_classes']
+        for key in required_data_keys:
+             if not OmegaConf.select(model_config, f'data.{key}', default=None):
+                 logger.error(f"Config missing required data field: data.{key}")
+                 return None, None, None
+        if not OmegaConf.select(model_config, 'id2color', default=None):
+             logger.error("Config missing required field: id2color")
+             return None, None, None
+        if not OmegaConf.select(model_config, 'training.model_name', default=None):
+             logger.error("Config missing required field: training.model_name")
+             return None, None, None
+
+        try:
+            # Convert OmegaConf structures to standard Python types using OmegaConf.to_container
+            # resolve=True handles variable interpolation (like ${data.base_dir}) if used in relevant fields
+            img_size = tuple(OmegaConf.to_container(model_config.data.image_size, resolve=True))
+            mean = list(OmegaConf.to_container(model_config.data.mean, resolve=True))
+            std = list(OmegaConf.to_container(model_config.data.std, resolve=True))
+            # Ensure keys in id2color are standard integers
+            id2color_map = {int(k): v for k, v in OmegaConf.to_container(model_config.id2color, resolve=True).items()}
+            num_classes = int(model_config.data.num_classes) # Ensure int
+
+            logger.debug(f"Converted config values: size={img_size}, mean={mean}, std={std}, id2color keys={list(id2color_map.keys())}, num_classes={num_classes}")
+
+            # Basic validation after conversion
+            if not isinstance(mean, list) or not isinstance(std, list) or not isinstance(id2color_map, dict): raise TypeError("Config values did not convert to list/dict.")
+            if len(mean) != 3 or len(std) != 3: raise ValueError(f"Incorrect mean/std length. Expected 3.") # Assuming 3 color channels
+            if len(img_size) != 2: raise ValueError(f"Incorrect image_size length. Expected 2 (H, W).")
+
+        except Exception as e:
+            logger.error(f"Error processing/converting configuration values: {e}")
+            logger.error(traceback.format_exc())
+            return None, None, None
+
+        # Define the image transformation pipeline
+        transform = transforms.Compose([
+            transforms.ToTensor(), # HWC [0,255] numpy -> CHW [0,1] torch
+            transforms.Resize(img_size, interpolation=transforms.InterpolationMode.BILINEAR), # Use converted tuple size, BILINEAR is common for images before model
+            transforms.Normalize(mean=mean, std=std) # Use converted lists
+        ])
+        logger.debug(f"Image transform applied for size {img_size}.")
+        input_tensor = transform(img_rgb).unsqueeze(0).to(device) # Add batch dim (B=1), move to device
+        logger.debug(f"Input tensor created with shape: {input_tensor.shape}") # Should be [1, 3, H, W]
+
+        # --- Run Prediction ---
+        with torch.no_grad():
+            logits = inference_model(input_tensor) # Expect [B, C, H, W] logits
+            logger.debug(f"Logits received with shape: {logits.shape}")
+            # Check logits shape again after potential upsampling in model forward
+            if logits.dim() != 4 or logits.shape[1] != num_classes:
+                 logger.error(f"Unexpected final logits shape or class number: {logits.shape}. Expected B x {num_classes} x H x W.")
+                 return None, None, None
+            # Argmax along class dimension (C), remove batch dim, move to CPU, convert type
+            pred_mask = logits.argmax(1).squeeze(0).cpu().numpy().astype(np.uint8) # H W, uint8
+            logger.debug(f"Prediction mask generated with shape: {pred_mask.shape}") # Should be [H, W]
+
+        # --- Post-processing ---
+        color_mask = num_to_rgb(pred_mask, id2color_map) # Use converted map
+        if color_mask is None: logger.error("num_to_rgb failed."); return None, None, None
+        logger.debug("Color mask generated.")
+
+        # Denormalize the *input tensor* for overlay display
+        denorm_img = denormalize(input_tensor, mean, std) # Use converted mean/std
+        if denorm_img is None: logger.error("denormalize failed."); return None, None, None
+        logger.debug("Input tensor denormalized for overlay.") # HWC, float32, [0,1]
+
+        # --- Create Overlay ---
+        # Ensure shapes match before blending (resize color mask to match denorm_img)
+        if denorm_img.shape[:2] != color_mask.shape[:2]:
+             logger.warning(f"Denorm img shape {denorm_img.shape[:2]} != Color mask shape {color_mask.shape[:2]}. Resizing color mask using INTER_NEAREST.")
+             # Resize color_mask (HWC float32) to match denorm_img (HWC float32)
+             color_mask = cv2.resize(color_mask, (denorm_img.shape[1], denorm_img.shape[0]), interpolation=cv2.INTER_NEAREST) # Use INTER_NEAREST for label masks
+
+        # Blend images: Original (denorm_img) * alpha + Mask (color_mask) * beta + gamma
+        overlay = cv2.addWeighted(denorm_img, 0.7, color_mask, 0.3, 0)
+        logger.debug("Overlay created using cv2.addWeighted.")
+        # overlay is HWC, float32, [0, 1], RGB
+
+        return denorm_img, color_mask, overlay
+
+    except Exception as e:
+        logger.error(f"Exception during inference pipeline for image: {e}")
+        logger.error(traceback.format_exc())
+        return None, None, None
+
+
+# --- Flask App ---
+app = Flask(__name__)
+CORS(app) # Allow all origins for API and Result routes resources={r"/api/*": {"origins": "*"}, r"/Result/*": {"origins": "*"}}
+logger.info("Flask app created and CORS enabled.")
+
+def allowed_file(filename):
+    """Checks if the filename has an allowed extension."""
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+# --- API Endpoints ---
+@app.route('/api/analyze', methods=['POST'])
+def analyze_image():
+    """Receives Base64 image, runs inference, saves original and overlay."""
+    global model, config, device # Access global vars
+    endpoint_log_prefix = "[POST /api/analyze]"
+    logger.info(f"{endpoint_log_prefix} Received request.")
+
+    # --- Basic Checks ---
+    if model is None: logger.error(f"{endpoint_log_prefix} Model not loaded."); return jsonify({"success": False, "message": "Model not loaded"}), 500
+    if not request.is_json: logger.warning(f"{endpoint_log_prefix} Not JSON."); return jsonify({"success": False, "message": "Request must be JSON"}), 400
+    data = request.get_json()
+    if not data or 'image' not in data or 'filename' not in data:
+        logger.warning(f"{endpoint_log_prefix} Missing image/filename in JSON body. Data received: {data}")
+        return jsonify({"success": False, "message": "Missing 'image' (base64) or 'filename' in JSON body"}), 400
+
+    base64_image_data = data['image']; original_filename = data['filename']
+    logger.info(f"{endpoint_log_prefix} Original filename from request: '{original_filename}'")
+    safe_original_filename = secure_filename(original_filename) # Sanitize
+    if not safe_original_filename or not allowed_file(safe_original_filename):
+        logger.warning(f"{endpoint_log_prefix} Invalid/disallowed filename after sanitization: '{safe_original_filename}' from '{original_filename}'")
+        return jsonify({"success": False, "message": "Invalid or disallowed filename/extension"}), 400
+    logger.info(f"{endpoint_log_prefix} Sanitized filename for saving/processing: '{safe_original_filename}'")
+
+    try:
+        # --- Decode Base64 ---
+        if ',' in base64_image_data: header, encoded = base64_image_data.split(',', 1)
+        else: encoded = base64_image_data # Assume no header
+        image_bytes = base64.b64decode(encoded)
+        logger.info(f"{endpoint_log_prefix} Base64 image decoded ({len(image_bytes)} bytes).")
+
+        # --- Save Original Image ---
+        original_path = os.path.join(UPLOAD_FOLDER, safe_original_filename)
+        try:
+            with open(original_path, "wb") as f: f.write(image_bytes)
+            logger.info(f"{endpoint_log_prefix} Original image saved to: '{original_path}'")
+        except Exception as e:
+            logger.error(f"{endpoint_log_prefix} Failed to save original image to '{original_path}': {e}")
+            return jsonify({"success": False, "message": "Failed to save uploaded image on server"}), 500
+
+        # --- Run Inference ---
+        logger.info(f"{endpoint_log_prefix} Starting inference for '{safe_original_filename}'...")
+        # Pass the global config object here
+        denorm_img, color_mask, overlay = run_inference_on_bytes(image_bytes, model, config, device)
+        if overlay is None: # Check if inference failed
+            logger.error(f"{endpoint_log_prefix} Inference pipeline returned None for '{safe_original_filename}'.")
+            return jsonify({"success": False, "message": "Inference process failed on server"}), 500
+        logger.info(f"{endpoint_log_prefix} Inference completed successfully for '{safe_original_filename}'.")
+
+        # --- Save Overlay Image ---
+        name_part, ext = os.path.splitext(safe_original_filename)
+        # Create consistent overlay filename (crucial for toggle endpoint)
+        overlay_filename = f"analyzed_{name_part}{ext}"
+        overlay_path = os.path.join(RESULT_FOLDER, overlay_filename)
+        logger.info(f"{endpoint_log_prefix} Determined overlay filename: '{overlay_filename}' -> path: '{overlay_path}'")
+
+        # Convert overlay (float32 HWC RGB [0,1]) to uint8 HWC BGR [0,255] for cv2.imwrite
+        try:
+            overlay_to_save_uint8 = (overlay * 255).astype(np.uint8)
+            overlay_to_save_bgr = cv2.cvtColor(overlay_to_save_uint8, cv2.COLOR_RGB2BGR)
+            save_success = cv2.imwrite(overlay_path, overlay_to_save_bgr)
+            if not save_success:
+                raise IOError(f"cv2.imwrite failed to save the overlay image to {overlay_path}")
+            logger.info(f"{endpoint_log_prefix} Overlay image saved successfully to: '{overlay_path}'")
+        except Exception as e:
+            logger.error(f"{endpoint_log_prefix} Failed to convert or save overlay image to '{overlay_path}': {e}")
+            logger.error(traceback.format_exc())
+            return jsonify({"success": False, "message": "Failed to save analysis result image"}), 500
+
+        # --- Success Response ---
+        logger.info(f"{endpoint_log_prefix} Analysis successful for '{safe_original_filename}'. Returning success.")
+        return jsonify({
+            "success": True,
+            "message": "Analysis complete",
+            # Optionally return relative paths for info, client mainly needs overlay_filename
+            "paths": {"original": os.path.relpath(original_path, BASE_DIR), "overlay": os.path.relpath(overlay_path, BASE_DIR)},
+            "overlay_filename": overlay_filename # Return the *exact* filename saved
+        }), 200
+
+    except base64.binascii.Error as e:
+        logger.error(f"{endpoint_log_prefix} Invalid Base64 data received: {e}")
+        return jsonify({"success": False, "message": "Invalid Base64 image data received"}), 400
+    except Exception as e:
+        logger.error(f"{endpoint_log_prefix} Unexpected error during analysis request processing: {e}")
+        logger.error(traceback.format_exc())
+        return jsonify({"success": False, "message": "Internal server error during analysis processing"}), 500
+
+
+@app.route('/api/toggle-image', methods=['GET'])
+def get_analysis_path():
+    """Checks if the analyzed version of a given original filename exists."""
+    endpoint_log_prefix = "[GET /api/toggle-image]"
+    logger.info(f"{endpoint_log_prefix} Received request.")
+    logger.info(f"{endpoint_log_prefix} Full request URL: {request.url}")
+    logger.info(f"{endpoint_log_prefix} Request Query Args: {request.args}") # Log received args
+
+    original_filename = request.args.get('filename') # Get filename from ?filename=...
+    if not original_filename:
+        logger.warning(f"{endpoint_log_prefix} Missing 'filename' query parameter.")
+        return jsonify({"message": "Missing 'filename' query parameter"}), 400
+
+    logger.info(f"{endpoint_log_prefix} Original filename received from query: '{original_filename}'")
+    safe_original_filename = secure_filename(original_filename) # Sanitize
+    if not safe_original_filename:
+         logger.warning(f"{endpoint_log_prefix} Invalid filename after sanitization: '{safe_original_filename}' from '{original_filename}'")
+         return jsonify({"message": "Invalid filename format"}), 400
+    logger.info(f"{endpoint_log_prefix} Sanitized filename for lookup: '{safe_original_filename}'")
+
+    # --- Construct Expected Overlay Path (MUST match /analyze logic) ---
+    name_part, ext = os.path.splitext(safe_original_filename)
+    expected_overlay_filename = f"analyzed_{name_part}{ext}"
+    expected_overlay_path = os.path.join(RESULT_FOLDER, expected_overlay_filename)
+    logger.info(f"{endpoint_log_prefix} Expecting overlay file at: '{expected_overlay_path}'")
+
+    # --- Check if File Exists ---
+    if os.path.exists(expected_overlay_path):
+        logger.info(f"{endpoint_log_prefix} Found analysis result file: '{expected_overlay_filename}'")
+        # Return just the filename, client constructs the full /Result/ URL
+        return jsonify({"filepath": expected_overlay_filename}), 200
+    else:
+        # Explicitly log the path that was checked and not found
+        logger.warning(f"{endpoint_log_prefix} Analysis result file NOT FOUND at checked path: '{expected_overlay_path}'")
+        # Return 404 Not Found status code
+        return jsonify({"message": f"Analysis result not found for '{original_filename}'"}), 404
+
+
+@app.route('/Result/<filename>')
+def serve_result_image(filename):
+    """Serves images from the RESULT_FOLDER."""
+    endpoint_log_prefix = "[GET /Result]"
+    # Sanitize filename received in URL path for security
+    safe_filename = secure_filename(filename)
+    if safe_filename != filename:
+         # Log if the requested filename was changed by sanitization
+         logger.warning(f"{endpoint_log_prefix} Requested filename '{filename}' was sanitized to '{safe_filename}'. Serving sanitized version.")
+
+    logger.info(f"{endpoint_log_prefix} Attempting to serve file: '{safe_filename}' from directory: '{RESULT_FOLDER}'")
+    try:
+        # Use Flask's send_from_directory - safer than manual path joining
+        # as_attachment=False means display in browser if possible
+        return send_from_directory(RESULT_FOLDER, safe_filename, as_attachment=False)
+    except FileNotFoundError:
+        # Log the specific file that was not found
+        logger.error(f"{endpoint_log_prefix} Requested file not found in result folder: '{safe_filename}'")
+        # Return 404 Not Found
+        return jsonify({"message": "Requested analysis image not found"}), 404
+    except Exception as e:
+        # Catch other potential errors (e.g., permission issues)
+        logger.error(f"{endpoint_log_prefix} Error serving file '{safe_filename}': {e}")
+        logger.error(traceback.format_exc())
+        # Return 500 Internal Server Error
+        return jsonify({"message": "Error serving analysis image"}), 500
+
+
+# --- Main Execution ---
+if __name__ == '__main__':
+    # Ensure model loaded successfully before starting server
+    if model:
+        logger.info("Model loaded successfully. Starting Flask development server...")
+        # Use debug=True for development (auto-reload, debugger)
+        # Use debug=False for production!
+        # host='0.0.0.0' makes it accessible on the network
+        app.run(host='0.0.0.0', port=7860, debug=True)
+    else:
+        # This message should appear if load_trained_model returned None
+        logger.critical("APPLICATION FAILED TO START: MODEL COULD NOT BE LOADED.")
+        # Exit code 1 indicates an error
+        exit(1)

config/config.yaml

@@ -0,0 +1,111 @@
+defaults:
+  - _self_
+
+data:
+  num_classes: 14
+  image_size: [512, 512]
+  mean: [0.4810, 0.4810, 0.4810]  #[0.4513, 0.4513, 0.4513]
+  std: [0.2492, 0.2492, 0.2492] #[0.1879, 0.1879, 0.1879]
+  background_cls_id: 0
+  split_ratio: 0.8
+  base_dir: 'E:\Seg\Test\medical_segmentation\Bitewing'
+  mask_dir: "${data.base_dir}/newAnnotations/"
+  original_dir: "${data.base_dir}/Images/"
+  dataset_path: "${hydra:runtime.cwd}/Bitwening_dataset1"
+  
+  # Normalization settings
+  normalization:
+    always_apply: true
+  
+  # Image resize configuration
+  resize:
+    interpolation: "INTER_NEAREST"  # CV2 interpolation method
+  
+  # Complete augmentation configuration with all parameters
+  augmentation:
+    # HorizontalFlip
+    use_horizontal_flip: true
+    horizontal_flip_prob: 0.5
+    
+    # VerticalFlip
+    use_vertical_flip: true
+    vertical_flip_prob: 0.5
+    
+    # ShiftScaleRotate
+    use_shift_scale_rotate: true
+    shift_scale_rotate_prob: 0.5
+    rotate_limit: 0.15
+    scale_limit: 0.12
+    shift_limit: 0.12
+    border_mode: 4  # cv2.BORDER_REFLECT_101
+    
+    # RandomBrightnessContrast
+    use_brightness_contrast: true
+    brightness_contrast_prob: 0.5
+    brightness_limit: 0.2
+    contrast_limit: 0.2
+    
+    # CoarseDropout
+    use_coarse_dropout: true
+    coarse_dropout:
+      max_holes: 8
+      min_holes: 5
+      max_height: 25
+      max_width: 25
+      fill_value: 0
+      mask_fill_value: 0
+      prob: 0.5
+
+training:
+  batch_size: 10
+  num_epochs: 1
+  init_lr: 3e-4
+  optimizer_name: "AdamW"
+  weight_decay: 0.1
+  use_scheduler: true
+  scheduler: "MultiStepLR"
+  model_name: "nvidia/segformer-b4-finetuned-ade-512-512"
+  num_workers: 0
+  pin_memory: true
+  drop_last: true
+  shuffle_train: true
+  shuffle_valid: false
+
+inference:
+  batch_size: 10
+  num_batches: 3
+
+wandb:
+  project: "UM_medical_segmentation"
+  log_model: true
+
+trainer:
+  accelerator: "gpu"
+  devices: "1"
+  strategy: "auto"
+  precision: "16-mixed"
+  enable_model_summary: false
+
+id2color:
+  0: [0, 0, 0]         # Black
+  1: [0, 0, 255]       # Blue
+  2: [0, 255, 0]       # Green
+  3: [255, 0, 0]       # Red
+  4: [255, 255, 0]     # Yellow
+  5: [255, 165, 0]     # Orange
+  6: [128, 0, 128]     # Purple
+  7: [0, 255, 255]     # Cyan
+  8: [255, 20, 147]    # Deep Pink
+  9: [75, 0, 130]      # Indigo
+  10: [139, 69, 19]    # Saddle Brown
+  11: [255, 192, 203]  # Pink
+  12: [47, 79, 79]     # Dark Slate Gray
+  13: [173, 255, 47]   # Green Yellow
+  14: [0, 128, 128]    # Teal
+
+
+
+experiment:
+  name: "EXPERIMENT_1 bitewing dataset"          # Descriptive name for the experiment
+  description: "Testing on 14 classes dataset with no layers and just fine tuning the model "    # What this experiment is testing
+  goal: "to check the performance of the model on 100 epochs"    # What you hope to achieve