🙋?♂?寫在最前面的

🥳博客主頁：😋 滑鼠滑輪不會動 😋
🏅歡迎關注🎯 點贊👍 收藏?? 留言📋
🎬首發時間：🎉2022年1月15日🎉
🆕最新更新時間：🎉2022年1月15日🎉
🌟本文由 滑鼠滑輪不會動 原創，CSDN首發！
📬不積跬步無以至千里！
📌本人水平非常有限，如有錯誤，請留言捶我小腦袋，萬分感謝感謝感謝！

熱門文章目錄
??跑通代碼—2020-CVPR–StegaStamp: Invisible Hyperlinks in Physical Photographs
??跑通代碼—2020-WACV-Dense Extreme Inception Network: Towards a Robust CNN Model for Edge Detection
??跑通代碼—2021-ICCV-HiNet: Deep Image Hiding by Invertible Network
??代碼復現—2021-IEEE-TETCI-A Generalized Deep Neural Network Approach for Digital Watermarking Analysis
??跑通代碼—2018-ECCV-HiDDeN: Hiding Data With Deep Networks
??跑通代碼—2018-CVPR—D-LinkNet: DeepGlobe-Road-Extraction-Challenge

在這里插入圖片描述

🙋?♂?寫在最前面的
?代碼整理初衷
?基礎設定
- 🌿必要的包
- 🌿基礎檢查與設定
- - 🎄訓練設備的檢測與選擇
  - 🎄引數的設定
  - 🎄路徑的檢查與創建
?資料集的制作和匯入
- - 🎄單類影像資料集匯入
  - 🎄帶有groundtruth的資料集匯入
?訓練程序資料和影像的及時反饋
?所有代碼的匯總
?后記

?代碼整理初衷

在復現演算法和看大佬們寫的代碼程序中發現，很多代碼都很模塊化，需要修改的地方無非是模型部分和一些前處理或者后處理的部分，其他構建資料集，設定引數都是很模塊化的代碼，根本不需要花費太多的時間，所以，想著可以將自己認為比較模塊化的代碼整理出來，方便自己和大家今后寫代碼直接查找和使用，需要對應部分代碼的可以復制部分，也可以在最后將所有的代碼一起復制，🤩

?基礎設定

🌿必要的包

import os  
import torch
import time  # 時間相關的包，訓練中計時
import numpy as np  # 數值計算
from glob import glob  # 獲取相匹配的路徑下檔案，讀取資料時使用
from PIL import Image, ImageOps  # 強大的影像處理包
from torch.utils.data import Dataset, DataLoader  # 批量匯入影像資料使用
from torch.utils.tensorboard import SummaryWriter  # 訓練程序中資料監視器
from torchvision import transforms  # 對影像資料進行預處理變換
import cv2  # 強大的影像處理包
import random  # 產生亂數

🌿基礎檢查與設定

🎄訓練設備的檢測與選擇

  device = torch.device('cpu' if torch.cuda.device_count() == 0 else 'cuda')

🎄引數的設定

🔶方法1??：直接使用 parser = argparse.ArgumentParser(), 使用 parser.add_argument() 增加設定，使用 args.batch_size 來呼叫

def parse_args():
    parser = argparse.ArgumentParser(description='DexiNed trainer.')
    parser.add_argument('--batch_size',type=int, default=8, help=''the mini-batch size (default: 8)'')
    parser.add_argument('--choose_test_data', type=int, default=0, help='Already set the dataset for testing choice: 0 - 8')
    parser.add_argument('--input_dir',type=str, default="dataset/BIPED", help='the path to the directory with the input data.')
    parser.add_argument('--log_path',type=str, default="logs/test0115", help='the path to the directory with the input data.')
    args = parser.parse_args()
    return args

🔶方法2??：在.yaml檔案單獨存放引數設定變數，使用 args.lr 來呼叫

with open('setting.yaml', 'r') as f:
    args = EasyDict(yaml.load(f, Loader=yaml.SafeLoader))


#  setting.yaml檔案中的設定如下，使用變數+：+引數的設定方法
cuda: True
exp_name: test0115
num_steps: 14000
batch_size: 16
lr: 0.0001
train_path: 'D:\data\mirflickr'
logs_path: "./logs/"
checkpoints_path: './checkpoints/'
saved_models: './saved_models'


# train函式中呼叫直接使用如下命令 args.lr
optimize_secret_loss = optim.Adam(g_vars, lr=args.lr)

🎄路徑的檢查與創建

🔶檢查路徑，如果不存在創建路徑; 拼接路徑

# 檢查路徑，如果不存在創建路徑
if not os.path.exists(args.checkpoints_path):
 	 os.makedirs(args.checkpoints_path)
 
# 拼接路徑
checkpoints_path = os.path.join(args.checkpoints_path, str(args.exp_name))

#創建路徑
os.makedirs(checkpoints_path, exist_ok=True)

?資料集的制作和匯入

🎄單類影像資料集匯入

class MakeData(Dataset):
    def __init__(self, data_path, size=(400, 400)):
        self.data_path = data_path
        self.size = size
         # .jpg根據自己的資料集進行替換，例如.png
        self.files_list = glob(os.path.join(self.data_path, '*.jpg'))
        self.to_tensor = transforms.ToTensor()

    def __getitem__(self, idx):    
        img_cover_path = self.files_list[idx]
        # 讀取影像并轉化成RGB格式
        img_cover = Image.open(img_cover_path).convert('RGB')
        # ImageOps.fit用于獲得制定裁剪尺寸的影像
        img_cover = ImageOps.fit(img_cover, self.size)
        img_cover = self.to_tensor(img_cover)
        return img_cover

    def __len__(self):
        return len(self.files_list)


# train函式中使用如下命令呼叫
dataset = MakeData(args.train_path, size=(400, 400))
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True)

🎄帶有groundtruth的資料集匯入

🔶以邊緣提取的資料集Biped為例，資料集分為原始影像和邊緣的ground truth影像

class BipedDataset(Dataset):
    def __init__(self,
                 data_root,
                 img_height,
                 img_width,
                 mean_bgr,
                 train_mode='train',
                 dataset_type='rgbr',
                 crop_img=False,
                 arg=None
                 ):
        self.data_root = data_root
        self.train_mode = train_mode
        self.dataset_type = dataset_type
        self.data_type = 'real'  # be aware that this might change in the future
        self.img_height = img_height
        self.img_width = img_width
        self.mean_bgr = mean_bgr
        self.crop_img = crop_img
        self.arg = arg
        self.data_index = self._build_index()
    def _build_index(self):
   		 # os.path.abspath取指定檔案或目錄的絕對路徑（完整路徑）
        data_root = os.path.abspath(self.data_root)
        sample_indices = []
        images_path = os.path.join(data_root, 'imgs', self.train_mode， self.dataset_type)
        labels_path = os.path.join(data_root, 'edge_maps', self.train_mode, self.dataset_type)
        for directory_name in os.listdir(images_path):
            image_directories = os.path.join(images_path, directory_name)
            for file_name_ext in os.listdir(image_directories):
                file_name = os.path.splitext(file_name_ext)[0]
                sample_indices.append(
                    (os.path.join(images_path, directory_name, file_name + '.jpg'),
                     os.path.join(labels_path, directory_name, file_name + '.png'),)
                ) # 這里的.jpg和.png是根據自己資料集更換的
        return sample_indices

    def __len__(self):
        return len(self.data_index)

    def __getitem__(self, idx):
        # get data sample
        image_path, label_path = self.data_index[idx]
        # 讀取影像資料
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
        image, label = self.transform(img=image, gt=label)
        return dict(images=image, labels=label)

    def transform(self, img, gt):
        gt = np.array(gt, dtype=np.float32)
        if len(gt.shape) == 3:
            gt = gt[:, :, 0]
        gt /= 255. # 歸一化
        img = np.array(img, dtype=np.float32)
        img -= self.mean_bgr  # 取均值處理
        i_h, i_w,_ = img.shape
        # 設定隨機剪裁的尺寸為400
        crop_size = self.img_height if self.img_height == self.img_width else 400# 
        # # for BIPED
        if np.random.random() > 0.5: 
             h,w = gt.shape
            LR_img_size = 256  
            i = random.randint(0, h - LR_img_size)
            j = random.randint(0, w - LR_img_size)
            # if img.
            img = img[i:i + LR_img_size , j:j + LR_img_size ]
            gt = gt[i:i + LR_img_size , j:j + LR_img_size ]
            img = cv2.resize(img, dsize=(crop_size, crop_size),)
            gt = cv2.resize(gt, dsize=(crop_size, crop_size))
        else:
            # 如果不隨機剪裁就重置影像尺寸
            img = cv2.resize(img, dsize=(crop_size, crop_size))
            gt = cv2.resize(gt, dsize=(crop_size, crop_size))
			  # 對影像進行變換，由原來的 H W C-> C H W
        img = img.transpose((2, 0, 1))
        img = torch.from_numpy(img.copy()).float() #  轉成tensor
        gt = torch.from_numpy(np.array([gt])).float()
        return img, gt

?訓練程序資料和影像的及時反饋

🔶使用tensorborad對訓練中的資料和影像進行展示

# 設定路徑保存資料和影像
writer = SummaryWriter(args.log_path)
# 增加要展示的資料，其中 loss/image_loss中的loss代表分組，image_loss是對應資料的名稱，global_step為全域的步數，告訴writer什么時候展示
writer.add_scalar('loss/image_loss', image_loss, global_step)
writer.add_scalar('loss/secret_loss', secret_loss, global_step)
# 增加不同的分組，設定metric的分組
writer.add_scalar('metric/bit_acc', bit_acc, global_step)
writer.add_scalar('metric/str_acc', str_acc, global_step)


# 增加影像的展示，設定了三個分組，input，encoded和transformed
writer.add_image('input/image_input', image_input[0], global_step)
writer.add_image('input/image_warped', input_warped[0], global_step)

writer.add_image('encoded/encoded_warped', encoded_warped[0], global_step)
writer.add_image('encoded/residual_warped', residual_warped[0] + 0.5, global_step)
writer.add_image('encoded/encoded_image', encoded_image[0], global_step)

writer.add_image('transformed/transformed_image', transformed_image[0], global_step)
writer.add_image('transformed/test', test_transform[0], global_step)

# 在global_step訓練完成之后關閉writer
writer.close()

# 訓練程序中在terminal中輸入對應的路徑，查看程序中的結果,其中tensorboard --logdir不變，只需要輸入保存的log_path即可，如果想退出 CTRL+C
PS D:\StegaStamp_pytorch_original> tensorboard --logdir logs/test0115

在terminal中輸入tensorboard --logdir logs/test0115之后，點擊出現的本地鏈接 http://localhost:6006/ ，會彈出tensorborad的監視網頁，能實時看到訓練的結果

在這里插入圖片描述

?所有代碼的匯總


# ---------------------------------基礎的包匯入----------------------------------

import os  
import torch
import time  # 時間相關的包，訓練中計時
import numpy as np  # 數值計算
from glob import glob  # 獲取相匹配的路徑下檔案，讀取資料時使用
from PIL import Image, ImageOps  # 強大的影像處理包
from torch.utils.data import Dataset, DataLoader  # 批量匯入影像資料使用
from torch.utils.tensorboard import SummaryWriter  # 訓練程序中資料監視器
from torchvision import transforms  # 對影像資料進行預處理變換
import cv2  # 強大的影像處理包
import random  # 產生亂數


# --------------------------------訓練設備的檢測與選擇--------------------------------
 device = torch.device('cpu' if torch.cuda.device_count() == 0 else 'cuda')
 
 # 引數的設定，使用 **parser.add_argument()** 增加設定，使用 **args.batch_size** 來呼叫
 def parse_args():
    parser = argparse.ArgumentParser(description='DexiNed trainer.')
    parser.add_argument('--batch_size',type=int, default=8, help=''the mini-batch size (default: 8)'')
    parser.add_argument('--choose_test_data', type=int, default=0, help='Already set the dataset for testing choice: 0 - 8')
    parser.add_argument('--input_dir',type=str, default="dataset/BIPED", help='the path to the directory with the input data.')
    args = parser.parse_args()
    return args
    

# --------------------------------路徑相關操作--------------------------------
#  檢查路徑，如果不存在創建路徑
if not os.path.exists(args.checkpoints_path):
 	os.makedirs(args.checkpoints_path)
 
# 拼接路徑
log_path = os.path.join(args.logs_path, str(args.exp_name))

#創建路徑
os.makedirs(checkpoints_path, exist_ok=True)



# -------------------------單類資料集的處理---------------------------------
class MakeData(Dataset):
    def __init__(self, data_path, size=(400, 400)):
        self.data_path = data_path
        self.size = size
         # .jpg根據自己的資料集進行替換，例如.png
        self.files_list = glob(os.path.join(self.data_path, '*.jpg'))
        self.to_tensor = transforms.ToTensor()

    def __getitem__(self, idx):
    
        img_cover_path = self.files_list[idx]
        # 讀取影像并轉化成RGB格式
        img_cover = Image.open(img_cover_path).convert('RGB')
        # ImageOps.fit用于獲得制定裁剪尺寸的影像
        img_cover = ImageOps.fit(img_cover, self.size)
        img_cover = self.to_tensor(img_cover)

        return img_cover, secret

    def __len__(self):
        return len(self.files_list)


# train函式中使用如下命令呼叫
dataset = MakeData(args.train_path, size=(400, 400))
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True)



# -------------------------帶有groundtruth的資料集匯入---------------------------------

class BipedDataset(Dataset):
    def __init__(self,
                 data_root,
                 img_height,
                 img_width,
                 mean_bgr,
                 train_mode='train',
                 dataset_type='rgbr',
                 crop_img=False,
                 arg=None
                 ):
        self.data_root = data_root
        self.train_mode = train_mode
        self.dataset_type = dataset_type
        self.data_type = 'real'  # be aware that this might change in the future
        self.img_height = img_height
        self.img_width = img_width
        self.mean_bgr = mean_bgr
        self.crop_img = crop_img
        self.arg = arg
        self.data_index = self._build_index()
    def _build_index(self):
   		 # os.path.abspath取指定檔案或目錄的絕對路徑（完整路徑）
        data_root = os.path.abspath(self.data_root)
        sample_indices = []
        images_path = os.path.join(data_root, 'imgs', self.train_mode， self.dataset_type)
        labels_path = os.path.join(data_root, 'edge_maps', self.train_mode, self.dataset_type)
        for directory_name in os.listdir(images_path):
            image_directories = os.path.join(images_path, directory_name)
            for file_name_ext in os.listdir(image_directories):
                file_name = os.path.splitext(file_name_ext)[0]
                sample_indices.append(
                    (os.path.join(images_path, directory_name, file_name + '.jpg'),
                     os.path.join(labels_path, directory_name, file_name + '.png'),)
                ) # 這里的.jpg和.png是根據自己資料集更換的
        return sample_indices

    def __len__(self):
        return len(self.data_index)

    def __getitem__(self, idx):
        # get data sample
        image_path, label_path = self.data_index[idx]
        # 讀取影像資料
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
        image, label = self.transform(img=image, gt=label)
        return dict(images=image, labels=label)

    def transform(self, img, gt):
        gt = np.array(gt, dtype=np.float32)
        if len(gt.shape) == 3:
            gt = gt[:, :, 0]
        gt /= 255. # 歸一化
        img = np.array(img, dtype=np.float32)
        img -= self.mean_bgr  # 取均值處理
        i_h, i_w,_ = img.shape
        # 設定隨機剪裁的尺寸為400
        crop_size = self.img_height if self.img_height == self.img_width else 400# 
        # # for BIPED
        if np.random.random() > 0.5: 
             h,w = gt.shape
            LR_img_size = 256  
            i = random.randint(0, h - LR_img_size)
            j = random.randint(0, w - LR_img_size)
            # if img.
            img = img[i:i + LR_img_size , j:j + LR_img_size ]
            gt = gt[i:i + LR_img_size , j:j + LR_img_size ]
            img = cv2.resize(img, dsize=(crop_size, crop_size),)
            gt = cv2.resize(gt, dsize=(crop_size, crop_size))
        else:
            # 如果不隨機剪裁就重置影像尺寸
            img = cv2.resize(img, dsize=(crop_size, crop_size))
            gt = cv2.resize(gt, dsize=(crop_size, crop_size))
			  # 對影像進行變換，由原來的 H W C-> C H W
        img = img.transpose((2, 0, 1))
        img = torch.from_numpy(img.copy()).float() #  轉成tensor
        gt = torch.from_numpy(np.array([gt])).float()
        return img, gt


# -------------------------訓練程序資料和影像的及時反饋---------------------------------

# 設定路徑保存資料和影像
writer = SummaryWriter(args.log_path)
# 增加要展示的資料，其中 **loss/image_loss** 中的loss代表分組，image_loss是對應資料的名稱，global_step為全域的步數，告訴writer什么時候展示
writer.add_scalar('loss/image_loss', image_loss, global_step)
writer.add_scalar('loss/secret_loss', secret_loss, global_step)
# 增加不同的分組，設定metric的分組
writer.add_scalar('metric/bit_acc', bit_acc, global_step)
writer.add_scalar('metric/str_acc', str_acc, global_step)


# 增加影像的展示，設定了三個分組，input，encoded和transformed
writer.add_image('input/image_input', image_input[0], global_step)
writer.add_image('input/image_warped', input_warped[0], global_step)

writer.add_image('encoded/encoded_warped', encoded_warped[0], global_step)
writer.add_image('encoded/residual_warped', residual_warped[0] + 0.5, global_step)
writer.add_image('encoded/encoded_image', encoded_image[0], global_step)

writer.add_image('transformed/transformed_image', transformed_image[0], global_step)
writer.add_image('transformed/test', test_transform[0], global_step)

# 在global_step訓練完成之后關閉writer
writer.close()

# 訓練程序中在terminal中輸入對應的路徑，查看程序中的結果,其中tensorboard --logdir不變，只需要輸入保存的log_path即可
PS D:\StegaStamp_pytorch_original> tensorboard --logdir logs/test0115

?后記

歸納了一些使用pytorch進行影像處理中經常使用的模塊化代碼，學習與實踐中還會不斷更新，歡迎大家在評論區留言自己遇到的好用的模塊化代碼，一起來完善，😃

轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/413657.html

標籤：其他

上一篇：[YOLO專題-27]：YOLO V5 小目標檢測遇到的問題與常見解決辦法

下一篇：【影像配準】基于粒子群改進的sift影像配準matlab原始碼

影像處理---Pytorch中模塊化的代碼整理（持續更新中......）

🙋?♂?寫在最前面的

目錄

?代碼整理初衷

?基礎設定

🌿必要的包

🌿基礎檢查與設定

🎄訓練設備的檢測與選擇

🎄引數的設定

🎄路徑的檢查與創建

?資料集的制作和匯入

🎄單類影像資料集匯入

🎄帶有groundtruth的資料集匯入

?訓練程序資料和影像的及時反饋

?所有代碼的匯總

?后記