【Python】一次元rgbリストと３次元rgb画像テンソルを相互変換する(PyTorch,Numpy)

プログラミング

2022.11.26

1 はじめに
2 画像からrgbリストを取り出す
3 終わりに

はじめに
画像からrgbリストを取り出す
終わりに

はじめに

どーも、可視化研究室でデータサイエンスを学んでいるゆうき(@engieerblog_Yu)です。

研究でGANを作成している際に、rgbリストと画像の変換が必要になったのでログを残しておきます。

画像からrgbリストを取り出す

まずはrgb画像から、以下ようなリストを取り出すということをしました。

\([r_0,g_0,b_0,r_1,g_1,b_1・・・,W,H]\)

それぞれの値は、ピクセルごとのrgb値です。

最後にrgb画像のサイズを追加します。

以下のコードで、torch.tensor型の画像テンソルを上記のようなリストに変えることができました。

#image:torch.tensor(dimension:3)

def generate_rgb_array(image):
  W = image.shape[1]
  H = image.shape[2]
  array = []
  for i in range(W):
    for j in range(H):
      for n in range(3):
        array.append(image[n][i][j].item())

  array.append(W)
  array.append(H)

  return array

次に先ほどのリストを、torch.tensor型の三次元画像テンソルにしていきます。

#rgb:list
#W,H:scaler

def generate_image(rgb,W,H):
    k = 0
    image = np.zeros((3,W,H))
    for w in range(W):
        for h in range(H):
            for i in range(3):
                image[i][w][h] = rgb[k]
                k += 1

    image = torch.from_numpy(image.astype(np.float32)).clone()
    return image

ちなみに画像を出力したい場合は、以下で簡単に出力することができます。