在Google Colab中访问带类别图像并用于FFT图像处理的方法

我在Google Colab中尝试打开图像文件，目标是将FFT应用于图像处理。目前通过代码读取到的是标注为数组值的文件，想了解如何访问图像本身而非数值。我的Drive文件夹中DATASET目录包含A、B、C、D等子文件夹，每个子文件夹下有大量字母图像，需读取这些带类别的图像，实现打印指定图像、将图像送入FFT算法处理后生成新数据集的需求。当前代码如下：

!nvidia-smi
!nvidia-smi
!nvidia-smi 
!pip install tensorflow-gpu
!pip install tensorflow_hub
from __future__ import absolute_import, division, print_function, unicode_literals
import matplotlib.pylab as plt
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
import pandas as pd
import os
import cv2

from google.colab import drive
drive.mount('/content/drive')
    
data_root='/content/drive/My Drive/ML/AlpDatabase/DATASET'
   
import numpy as np
from matplotlib import pyplot as plt
    
data = tf.keras.utils.image_dataset_from_directory('/content/drive/My Drive/ML/AlpDatabase/DATASET')
    

data_iterator = data.as_numpy_iterator()
    
batch = data_iterator.next()
    
batch[0].shape
    
fig, ax = plt.subplots(ncols=4, figsize=(20,20))
    for idx, img in enumerate(batch[0][:4]):
      ax[idx].imshow(img.astype(int))
      ax[idx].title.set_text(batch[1][idx])

1. 图像与数组的关系说明 #

你看到的数组就是图像本身——图像在代码中本质就是像素值组成的矩阵，plt.imshow()可以将这个数组还原为可视化的图像，二者是同一数据的不同表现形式。

2. 打印指定图像 #

通过数据集迭代器定位目标图像，示例代码如下：

# 重置数据集迭代器
data_iterator = data.as_numpy_iterator()

# 假设目标图像在第2个批次（索引从0开始，这里遍历1次跳过第一个批次）
for _ in range(1):
    batch = data_iterator.next()

# 获取该批次中第3张图像（索引从0开始）
target_img = batch[0][2]
target_label = batch[1][2]

# 可视化指定图像
plt.figure(figsize=(8,8))
plt.imshow(target_img.astype(int))
plt.title(f"标签: {target_label}")
plt.axis('off')
plt.show()

3. FFT处理图像并生成新数据集 #

定义FFT处理函数，遍历整个数据集完成转换，同时保留标签并保存结果：

import numpy as np
import cv2

def apply_fft(img):
    # 转换为灰度图（适配RGB图像）
    if img.shape[-1] == 3:
        gray = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2GRAY)
    else:
        gray = img.astype(np.uint8)
    # 执行FFT并移频到中心
    fft = np.fft.fft2(gray)
    fft_shift = np.fft.fftshift(fft)
    # 计算幅度谱并归一化到0-255范围
    magnitude = 20 * np.log(np.abs(fft_shift))
    magnitude = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    return magnitude

# 遍历数据集生成FFT处理后的结果
fft_images = []
labels = []
for batch in data.as_numpy_iterator():
    imgs, lbls = batch
    for img, lbl in zip(imgs, lbls):
        fft_img = apply_fft(img)
        fft_images.append(fft_img)
        labels.append(lbl)

# 转换为数组格式并保存到Drive
fft_images = np.array(fft_images)
labels = np.array(labels)
np.save('/content/drive/My Drive/ML/AlpDatabase/FFT_dataset/images.npy', fft_images)
np.save('/content/drive/My Drive/ML/AlpDatabase/FFT_dataset/labels.npy', labels)

# 对比可视化原图像与FFT处理结果
plt.figure(figsize=(10,5))
plt.subplot(121)
plt.imshow(imgs[0].astype(int))
plt.title("原图像")
plt.axis('off')
plt.subplot(122)
plt.imshow(fft_images[0], cmap='gray')
plt.title("FFT幅度谱")
plt.axis('off')
plt.show()

4. 原代码的修正优化 #

原代码存在缩进错误，同时可增加图像大小统一等配置，修正后完整代码：

!nvidia-smi
!pip install tensorflow-gpu tensorflow_hub
from __future__ import absolute_import, division, print_function, unicode_literals
import matplotlib.pylab as plt
import tensorflow as tf
import numpy as np
import cv2

from google.colab import drive
drive.mount('/content/drive')

data_root='/content/drive/My Drive/ML/AlpDatabase/DATASET'

# 加载数据集，统一图像大小并设置批次
data = tf.keras.utils.image_dataset_from_directory(
    data_root,
    image_size=(256, 256),
    batch_size=32
)

data_iterator = data.as_numpy_iterator()
batch = data_iterator.next()

# 可视化批次图像
fig, ax = plt.subplots(ncols=4, figsize=(20,20))
for idx, img in enumerate(batch[0][:4]):
    ax[idx].imshow(img.astype(int))
    ax[idx].title.set_text(batch[1][idx])
    ax[idx].axis('off')
plt.show()

内容的提问来源于stack exchange，提问作者EmployedSlaveStudent