from tqdm import tqdm from torch import nn from sklearn.model_selection import KFold from skimage import measure from all_models_tools.all_model_tools import call_back from utils.Stomach_Config import Training_Config, Loading_Config, Save_Result_File_Config from Load_process.LoadData import Loding_Data_Root from Training_Tools.PreProcess import Training_Precesses from ..Models.GastroSegNet_Model import GastroSegNet from Model_Loss.Segmentation_Loss import Segmentation_Loss from draw_tools.draw import plot_history import time import torch.optim as optim import torch import torch.nn.functional as F import numpy as np import cv2 import os class Segmentation_Block_Training_Step(Loding_Data_Root, Training_Precesses): def __init__(self, Best_Model_Save_Root): self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 設置裝置，若有GPU則使用GPU，否則使用CPU self.Model = self.Construct_Segment_Model_CUDA() # 模型變數 self.train_subset = None # Training Dataset 的子集 self.val_subset = None # Validation Dataset 的子集 self.train_loader = None # Training DataLoader 的讀檔器 self.val_loader = None # Validation DataLoader 的讀檔器 self.Mask = None # 遮罩變數，接收RD Net產出來的Mask self.Grad = None # 梯度變數，後面用來執行Grad CAM self.model_name = Training_Config["Mask_Experiment_Name"] # 取名，使用哪個模型(可能是預處理模型/自己設計的模型) self.epoch = Training_Config["Epoch"] # 訓練該模型的次數 self.train_batch_size = Training_Config["Train_Batch_Size"] # 訓練模型的Batch Size self.Experiment_Name = Training_Config["Mask_Experiment_Name"] # 取名，使用哪個模型(可能是預處理模型/自己設計的模型) self.Best_Model_Save_Root = Best_Model_Save_Root # 初始化繼承物件 Training_Precesses.__init__(self, Training_Config["Image_Size"]) Loding_Data_Root.__init__(self, Loading_Config["Training_Labels"], Loading_Config["Train_Data_Root"], Loading_Config["Test_Data_Root"]) def Processing_Main(self, training_dataset, return_processed_images=False, test_dataloader=None): Best_Model_Path = None Best_Validation_Loss = 100000000 # K-Fold loop kf = KFold(n_splits=5, shuffle=True, random_state=42) Training_Data_Lader = self.Dataloader_Sampler(training_dataset, self.train_batch_size, True) for fold, (train_idx, val_idx) in enumerate(kf.split(range(len(training_dataset)))): # K-Fold 交叉驗證迴圈 print(f"\nStarting Fold {fold + 1}/5") # Create training and validation subsets for this fold self.train_subset = torch.utils.data.Subset(training_dataset, train_idx) self.val_subset = torch.utils.data.Subset(training_dataset, val_idx) # Wrap subsets in DataLoaders (use same batch size as original) self.train_loader = self.Dataloader_Sampler(self.train_subset , self.train_batch_size, True) self.val_loader = self.Dataloader_Sampler(self.val_subset, self.train_batch_size, True) # 模型訓練與驗證 model_path, Training_Losses, Validation_Losses, Total_Epoch, best_val_loss = self.Training_And_Validation(fold) Losses = [Training_Losses, Validation_Losses] if best_val_loss < Best_Validation_Loss: Best_Validation_Loss = best_val_loss Best_Model_Path = model_path # 將訓練結果化成圖，並將化出來的圖丟出去儲存 plot_history(Losses, None, f"{Save_Result_File_Config['Segument_Plot_Image']}/{self.Experiment_Name}", f"train-{str(fold)}") # 將訓練結果化成圖，並將化出來的圖丟出去儲存 # 如果需要返回處理後的圖像（用於後續識別訓練） if return_processed_images is not None: # 載入最佳模型 self.Model = self.Construct_Segment_Model_CUDA() self.Model.load_state_dict(torch.load(Best_Model_Path)) self.Model.eval() # 處理測試數據 with torch.no_grad(): for Input_Images, Mask_Ground_Truth_Image, Labels, File_Name, File_Classes in Training_Data_Lader: # 使用Model_Branch處理圖像並獲取分割結果，同時傳遞文件名以保存邊界框圖像 self.Model_Branch(Input_Images, Mask_Ground_Truth_Image, 0.0, Save_Result_File_Config["Segument_Bounding_Box_Image"], return_processed_image=True, file_names=File_Name, Classes=File_Classes) avg_test_loss = self.evaluate_on_test(Best_Model_Path, test_dataloader) return Best_Model_Path, avg_test_loss return Best_Model_Path def Training_And_Validation(self, Fold): self.Model = self.Construct_Segment_Model_CUDA() # 模型初始化 Optimizer = optim.SGD(self.Model.parameters(), lr=0.045, momentum=0.9, weight_decay=0.01) # 優化器初始化 model_path, early_stopping, scheduler = call_back(self.Best_Model_Save_Root, f"fold{Fold}", Optimizer) # 防止過擬合細節函數初始化 epoch = 0 Training_Losses, Validation_Losses = [], [] Training_Running_Losses, Validation_Running_Losses = 0.0, 0.0 # Epoch loop for epoch in range(self.epoch): self.Model.train() # Start training # Progress bar for training batches epoch_iterator = tqdm(self.train_loader, desc=f"Fold {Fold + 1}/5, Epoch [{epoch + 1}/{self.epoch}]") Start_Time = time.time() for Input_Images, Mask_Ground_Truth_Image, Labels, File_Name, File_Classes in epoch_iterator: Optimizer.zero_grad() # 清零梯度，防止梯度累積 Training_Total_Losses, Training_Running_Losses = self.Model_Branch(Input_Images, Mask_Ground_Truth_Image, Training_Running_Losses) Training_Total_Losses.backward() Optimizer.step() epoch_iterator = self.Calculate_Progress_And_Timing(Input_Images, self.train_subset, Training_Total_Losses, epoch_iterator, Start_Time) Training_Losses, Training_Running_Losses = self.Calculate_Average_Scores(self.train_loader, Training_Running_Losses, Training_Losses) # Validation step self.Model.eval() epoch_iterator_Validation = tqdm(self.val_loader, desc=f"\tValidation-Fold {Fold + 1}/5, Epoch [{epoch + 1}/{self.epoch}]") with torch.no_grad(): for Input_Images, Mask_Ground_Truth_Image, Labels, File_Name, File_Classes in epoch_iterator_Validation: Validation_Total_Losses, Validation_Running_Losses = self.Model_Branch(Input_Images, Mask_Ground_Truth_Image, Validation_Running_Losses) # 添加Start_Time參數 Start_Time = time.time() epoch_iterator_Validation = self.Calculate_Progress_And_Timing(Input_Images, self.val_subset, Validation_Total_Losses, epoch_iterator_Validation, Start_Time) Validation_Losses, Validation_Running_Losses = self.Calculate_Average_Scores(self.val_loader, Validation_Running_Losses, Validation_Losses) print(f"Traini Loss: {Training_Running_Losses:.4f}, Validation Loss: {Validation_Running_Losses:.4f}\n") # Early stopping early_stopping(Validation_Running_Losses, self.Model, model_path) if early_stopping.early_stop: print(f"Early stopping triggered in Fold {Fold + 1} at epoch {epoch + 1}") break # Scheduler step scheduler.step(Validation_Running_Losses) Total_Epoch = epoch + 1 best_val_loss = early_stopping.best_loss return model_path, Training_Losses, Validation_Losses, Total_Epoch, best_val_loss def Construct_Segment_Model_CUDA(self): GaSeg = GastroSegNet() # 添加輸出模型摘要的功能 print("\n==== GastroSegNet 模型摘要 ====\n") print(f"輸入通道數: {GaSeg.encoder[0].conv[0].in_channels}") print(f"輸出通道數: {GaSeg.final_conv.out_channels}") # 計算總參數量 total_params = sum(p.numel() for p in GaSeg.parameters() if p.requires_grad) print(f"可訓練參數總量: {total_params:,}") # 顯示模型結構 print("\n模型結構:") print(f"- 編碼器層數: {len(GaSeg.encoder)}") print(f"- 解碼器層數: {len(GaSeg.decoder)}") print("\n特徵通道配置:") features_str = ", ".join([str(GaSeg.encoder[i].conv[0].out_channels) for i in range(len(GaSeg.encoder))]) print(f" - 編碼器特徵通道: {features_str}") print(f" - 瓶頸層特徵通道: {GaSeg.bottleneck.conv[0].out_channels}") print("\n==== 摘要結束 ====\n") return self.Convert_Model_To_CUDA(GaSeg) def Convert_Model_To_CUDA(self, model): if torch.cuda.device_count() > 1: model = nn.DataParallel(model) model = model.to(self.device) return model def Model_Branch(self, Input_Images, Mask_Ground_Truth_Image, running_loss, Save_Dir = None, return_processed_image=False, file_names=None, Classes=None): # 直接將張量移到設備上，不需要重新創建 Input_Images.requires_grad = False Input_Images = Input_Images.to(self.device) Segmentation_Output = self.Model(Input_Images) # 如果需要返回處理後的圖像（用於推理階段） if return_processed_image: # 調整模型產出影像大小 Segmentation_Output = self.Compare_Image_And_Resize_It(Segmentation_Output, Input_Images) # 處理分割輸出，選擇候選框並將框外像素變黑，同時保存邊界框圖像 return self.process_segmentation_output(Input_Images, Segmentation_Output, Save_Dir, save_bbox_images=True, file_names=file_names, Classes=Classes) Mask_Ground_Truth_Image = Mask_Ground_Truth_Image.to(self.device) # 調整模型產出影像大小 Segmentation_Output = self.Compare_Image_And_Resize_It(Segmentation_Output, Mask_Ground_Truth_Image) Losses = self.Losses(Segmentation_Output, Mask_Ground_Truth_Image) # 計算損失（不需要手動設置requires_grad，因為損失計算會自動處理梯度） running_loss += Losses.item() return Losses, running_loss def Compare_Image_And_Resize_It(self, Image, Target): # 調整兩張影像大小到一樣 # 檢查Target的維度 if Target.dim() < 3: # 如果Target是2D張量，將其擴展為4D張量 [batch_size, channels, height, width] Target = Target.unsqueeze(0).unsqueeze(0) elif Target.dim() == 3: # 如果Target是3D張量，將其擴展為4D張量 [batch_size, channels, height, width] Target = Target.unsqueeze(0) # 獲取目標尺寸 target_height = Target.size(-2) # 使用倒數第二維作為高度 target_width = Target.size(-1) # 使用倒數第一維作為寬度 # 調整Image大小 Image = torch.nn.functional.interpolate(Image, size=(target_height, target_width), mode='nearest') # 動態調整通道維度 if Image.size(1) != Target.size(1) and Target.dim() >= 3: conv = torch.nn.Conv2d(Image.size(1), Target.size(1), kernel_size=1).to(self.device) Image = conv(Image) return Image def Losses(self, Segmentation_Output_Image, Segmentation_Mask_GroundTruth_Image): criterion = Segmentation_Loss() Loss = criterion(Segmentation_Output_Image, Segmentation_Mask_GroundTruth_Image) return Loss def Record_Average_Losses(self, DataLoader): loss = 0.0 losses = [] # Calculate average validation loss loss /= len(DataLoader) losses.append(loss) return losses def Calculate_Progress_And_Timing(self, inputs, Subset, loss, epoch_iterator, Start_Time): # Calculate progress and timing total_samples = len(Subset) processed_samples = 0 processed_samples += inputs.size(0) # Use size(0) for batch size # Calculate progress and timing elapsed_time = time.time() - Start_Time iterations_per_second = processed_samples / elapsed_time if elapsed_time > 0 else 0 eta = (total_samples - processed_samples) / iterations_per_second if iterations_per_second > 0 else 0 time_str = f"{int(elapsed_time//60):02d}:{int(elapsed_time%60):02d}<{int(eta//60):02d}:{int(eta%60):02d}" # Calculate batch metrics using PSNR/SSIM loss batch_loss = loss.item() # Update progress bar with PSNR/SSIM loss epoch_iterator.set_postfix_str( f"{processed_samples}/{total_samples} [{time_str}, {iterations_per_second:.2f}it/s, " f"loss={batch_loss:.3f}]" ) return epoch_iterator def Calculate_Average_Scores(self, Data_Loader, Running_Losses, Losses): Running_Losses /= len(Data_Loader) Losses.append(Running_Losses) return Losses, Running_Losses def process_segmentation_output(self, input_images, segmentation_output, bbox_save_dir = None, save_bbox_images=True, file_names=None, Classes = None): """處理分割輸出，選擇候選框並將框外像素變黑，同時保存邊界框圖像 Args: input_images: 原始輸入圖像 [B, C, H, W] segmentation_output: 分割模型輸出 [B, 1, H, W] save_bbox_images: 是否保存邊界框圖像 file_names: 圖像文件名列表，用於保存邊界框圖像 Returns: processed_images: 處理後的圖像，框外像素變黑 [B, 3, H, W] (始終確保輸出是3通道) """ # 將輸出轉換為二值掩碼 (閾值為0.5) Batch_Size_Of_Image = segmentation_output.size(0) binary_masks = (torch.sigmoid(segmentation_output) > 0.5).float() # 創建一個輸出張量，確保始終是3通道 # 獲取批次大小、高度和寬度 Batch_Size_Of_Image, _, height, width = input_images.size() # 創建3通道的輸出張量，不管輸入是幾個通道，輸出都是3通道 processed_images = torch.zeros(Batch_Size_Of_Image, 3, height, width, device=input_images.device) # 對批次中的每張圖像進行處理 for Batch_Size in range(Batch_Size_Of_Image): # 創建保存邊界框圖像的目錄 new_bbox_save_dir = bbox_save_dir new_bbox_save_dir = os.path.join(new_bbox_save_dir, Classes[Batch_Size]) if save_bbox_images and not os.path.exists(new_bbox_save_dir): os.makedirs(new_bbox_save_dir, exist_ok=True) print(f"創建邊界框圖像保存目錄: {new_bbox_save_dir}") # 獲取當前圖像的二值掩碼並轉換為numpy數組 mask = binary_masks[Batch_Size, 0].cpu().numpy().astype(np.uint8) # 使用連通區域分析找出所有候選區域 labeled_mask, num_labels = measure.label(mask, return_num=True, connectivity=2) if num_labels > 0: # 計算每個區域的面積 regions = measure.regionprops(labeled_mask) # 根據面積排序區域（從大到小） regions.sort(key=lambda x: x.area, reverse=True) # 選擇最大的區域作為最終掩碼 if len(regions) > 0: # 創建一個新的掩碼，只包含最大的區域 final_mask = np.zeros_like(mask) for coords in regions[0].coords: final_mask[coords[0], coords[1]] = 1 # 獲取最大區域的邊界框 bbox = regions[0].bbox # (min_row, min_col, max_row, max_col) # 將最終掩碼轉換回PyTorch張量 final_mask_tensor = torch.from_numpy(final_mask).float().to(self.device) # 確保掩碼與輸入圖像的尺寸匹配 if final_mask_tensor.shape != input_images[Batch_Size, 0].shape: # 調整掩碼大小以匹配輸入圖像 final_mask_tensor = torch.nn.functional.interpolate( final_mask_tensor.unsqueeze(0).unsqueeze(0), size=input_images[Batch_Size, 0].shape, mode='nearest' ).squeeze(0).squeeze(0) # 將掩碼應用到原始圖像上（保留框內像素，將框外像素變黑） # 處理不同通道數的輸入圖像 if input_images.size(1) == 1: # 單通道輸入 # 將單通道複製到3個通道 for Channel in range(3): processed_images[Batch_Size, Channel] = input_images[Batch_Size, 0] * final_mask_tensor elif input_images.size(1) == 3: # 三通道輸入 # 直接複製三個通道 for Channel in range(3): processed_images[Batch_Size, Channel] = input_images[Batch_Size, Channel] * final_mask_tensor else: # 其他通道數（不太可能，但為了健壯性） # 取前三個通道或複製第一個通道 for Channel in range(3): if Channel < input_images.size(1): processed_images[Batch_Size, Channel] = input_images[Batch_Size, Channel] * final_mask_tensor else: processed_images[Batch_Size, Channel] = input_images[Batch_Size, 0] * final_mask_tensor # 保存帶有邊界框的圖像 if save_bbox_images: # 將輸入圖像轉換為numpy數組並調整為適合顯示的格式 img_tensor = input_images[Batch_Size].clone().detach().cpu() img_np = img_tensor.permute(1, 2, 0).numpy() # 將圖像從[0,1]範圍轉換為[0,255]範圍 img_np = (img_np * 255).astype(np.uint8) # 確保圖像是連續的內存塊 img_np = np.ascontiguousarray(img_np) # 如果圖像是單通道的，轉換為三通道 if img_np.shape[2] == 1: img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2BGR) elif img_np.shape[2] == 3: # 確保是BGR格式（OpenCV默認格式） img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) # 繪製邊界框 min_row, min_col, max_row, max_col = bbox cv2.rectangle(img_np, (min_col, min_row), (max_col, max_row), (0, 255, 0), 2) # 生成保存文件名 if file_names is not None and Batch_Size < len(file_names): # 使用提供的文件名 file_name = os.path.basename(file_names[Batch_Size]) save_path = os.path.join(new_bbox_save_dir, f"bbox_{file_name}.png") else: # 使用索引作為文件名 save_path = os.path.join(new_bbox_save_dir, file_names) # 保存圖像 (已經是BGR格式，直接保存) cv2.imwrite(save_path, img_np) print(f"已保存邊界框圖像: {save_path}") else: # 如果沒有找到區域，則保留原始圖像，但確保是3通道 if input_images.size(1) == 1: # 單通道輸入 # 將單通道複製到3個通道 for Channel in range(3): processed_images[Batch_Size, Channel] = input_images[Batch_Size, 0] elif input_images.size(1) == 3: # 三通道輸入 # 直接複製三個通道 processed_images[Batch_Size] = input_images[Batch_Size] else: # 其他通道數 # 取前三個通道或複製第一個通道 for Channel in range(3): if Channel < input_images.size(1): processed_images[Batch_Size, Channel] = input_images[Batch_Size, Channel] else: processed_images[Batch_Size, Channel] = input_images[Batch_Size, 0] else: # 如果沒有找到任何區域，則保留原始圖像，但確保是3通道 if input_images.size(1) == 1: # 單通道輸入 # 將單通道複製到3個通道 for Channel in range(3): processed_images[Batch_Size, Channel] = input_images[Batch_Size, 0] elif input_images.size(1) == 3: # 三通道輸入 # 直接複製三個通道 for Channel in range(3): processed_images[Batch_Size, Channel] = input_images[Batch_Size, Channel] else: # 其他通道數 # 取前三個通道或複製第一個通道 for Channel in range(3): if Channel < input_images.size(1): processed_images[Batch_Size, Channel] = input_images[Batch_Size, Channel] else: processed_images[Batch_Size, Channel] = input_images[Batch_Size, 0] # 在沒有候選區域的情況下也保存處理後的圖像 if save_bbox_images: # 轉換為可保存的numpy格式（與上方保存邊界框圖像一致的流程） img_tensor = processed_images[Batch_Size].clone().detach().cpu() img_np = img_tensor.permute(1, 2, 0).numpy() img_np = (img_np * 255).astype(np.uint8) img_np = np.ascontiguousarray(img_np) # 保證三通道BGR格式 if img_np.shape[2] == 1: img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2BGR) elif img_np.shape[2] == 3: img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) # 生成保存文件名（使用提供的檔名或索引），並標記為無邊界框 if file_names is not None and isinstance(file_names, (list, tuple)) and Batch_Size < len(file_names): file_name = os.path.basename(file_names[Batch_Size]) save_path = os.path.join(new_bbox_save_dir, f"no_bbox_{file_name}.png") else: base_name = file_names if isinstance(file_names, str) and len(file_names) > 0 else f"no_bbox_{Batch_Size}.png" save_path = os.path.join(new_bbox_save_dir, base_name) cv2.imwrite(save_path, img_np) print(f"已保存無邊界框圖像: {save_path}") def evaluate_on_test(self, model_path, test_dataloader): if test_dataloader is None: raise ValueError("Test dataloader is required for evaluation.") self.Model = self.Construct_Segment_Model_CUDA() self.Model.load_state_dict(torch.load(model_path)) self.Model.eval() Test_Losses = [] test_loss = 0.0 with torch.no_grad(): for Input_Images, Mask_Ground_Truth_Image, _, _, _ in test_dataloader: losses, test_loss = self.Model_Branch(Input_Images, Mask_Ground_Truth_Image, test_loss) losses, test_loss = self.Calculate_Average_Scores(test_dataloader, test_loss, Test_Losses) print(f"Average Test Loss: {test_loss}") return test_loss