아주 자세한 논문 리뷰는 다음 게시물 참고 바람
https://minsunstudio.tistory.com/66
[떠먹여주는 논문 리뷰] TMPI : Tiled Multiplane Images for Practical 3D Photography(ICCV 23.10)
ICCV 2023Metahttps://arxiv.org/abs/2309.14291 Tiled Multiplane Images for Practical 3D PhotographyThe task of synthesizing novel views from a single image has useful applications in virtual reality and mobile computing, and a number of approaches to the p
minsunstudio.tistory.com
수정 업뎃중
TMPI/run.py
단일 이미지에서 3D 효과가 있는 새로운 시점의 영상을 생성
1️⃣ 입력 이미지를 받아 깊이 정보를 추출
2️⃣ 깊이 정보를 이용해 새로운 시점에서 본 듯한 이미지 생성
3️⃣ 여러 프레임을 합쳐 3D 효과의 영상(mp4)으로 저장
4️⃣ OpenGL을 사용하면 빠르게, PyTorch 렌더러를 사용하면 안정적으로 실행
👉 주요 기능
• 이미지에서 깊이 정보(Depth Map) 추출
• 이를 활용해 새로운 시점에서 본 듯한 이미지 생성
• 여러 프레임을 조합해 3D 효과가 있는 영상(mp4) 생성
import os
import torch
from dataset import TMPIDataset
from torch.utils.data import DataLoader
import numpy as np
import argparse
from tmpi import TMPI
import config
import math
from moviepy.editor import ImageSequenceClip
from tmpi_renderer_gl import TMPIRendererGL
from tmpi_renderer import TMPIRenderer✔ TMPIDataset: 이미지 데이터를 불러오는 클래스
✔ TMPI: TMPI 모델 (3D 변환을 담당)
✔ moviepy.editor: 여러 이미지를 하나의 영상(mp4)으로 변환
🔹 2) 3D 변환 함수 (render_3d_photo)
def render_3d_photo(model, sample, poses, use_gl=True):
src_rgb_tiles = sample['src_rgb_tiles'].cuda() # (b, n, 3, h, w)
src_disp_tiles = sample['src_disp_tiles'].cuda() # (b, n, 1, h, w)
src_rgb = sample['src_rgb'].cuda() # (b, 3, h, w)
src_disp = sample['src_disp'].cuda() # (b, 1, h, w)
cam_int = sample['cam_int']
sx = sample['sx']
sy = sample['sy']
✔ 입력 이미지(src_rgb_tiles)와 깊이 정보(src_disp_tiles)를 **GPU(cuda)**로 보냄.
✔ cam_int, sx, sy → 카메라 보정 정보
✔ poses → 새로운 시점을 만들기 위한 카메라 이동 정보
📍 3D 변환 실행
# Predict Tiled Multiplane Images
mpis, mpi_disp = model( src_rgb_tiles, src_disp_tiles, src_rgb, src_disp)
torch.cuda.empty_cache()
✔ TMPI 모델이 실행되어 Multiplane Image(MPI) 데이터를 생성
✔ MPI란? 여러 개의 반투명한 레이어를 쌓아 3D 효과를 내는 방식
📍 OpenGL을 사용할지 PyTorch Renderer를 사용할지 선택
h, w = src_rgb.shape[-2:]
if use_gl:
renderer = TMPIRendererGL(h, w)
tgt_rgb_syn = renderer( mpis.cpu(), mpi_disp.cpu(), poses, cam_int, sx, sy)
else:
print("\033[91mNot Using OpenGL. Rendering will be slow.\033[0m")
device = 'cpu'
renderer = TMPIRenderer(h, w, device)
K = sample['K']
tgt_rgb_syn = renderer( mpis.to(device),
mpi_disp.to(device),
poses.to(device).flip(dims=[0]),
K.to(device),
sx.to(device),
sy.to(device))
return tgt_rgb_syn
✔ OpenGL을 사용할 경우 TMPIRendererGL을 사용하여 빠르게 3D 렌더링
✔ OpenGL이 없으면 TMPIRenderer(PyTorch 기반)를 사용하지만 속도가 느림
🔹 3) 새로운 시점(Novel View) 이동 경로 만들기
def render_path(num_frames=90, r_x=0.03, r_y=0.03, r_z=0.03):
t = torch.arange(num_frames) / (num_frames - 1)
poses = torch.eye(4).repeat(num_frames, 1, 1)
poses[:, 0, 3] = r_x * torch.sin(2. * math.pi * t)
poses[:, 1, 3] = r_y * torch.cos(2. * math.pi * t)
poses[:, 2, 3] = r_z * (torch.cos(2. * math.pi * t) - 1.)
return poses
✔ poses는 새로운 시점을 생성하기 위한 카메라 이동 경로
✔ num_frames=90 → 90개의 프레임을 생성하여 영상(mp4)으로 만들 예정
✔ 카메라가 x, y, z 방향으로 조금씩 이동하며 부드러운 3D 효과를 만듦
🔹 4) 전체 파이프라인 (test 함수)
def test(indir, outdir, use_gl_renderer, crop_edges=True):
data = TMPIDataset(data_root=indir)
dataset_size = len(data)
dataset_loader = DataLoader(
dataset=data,
batch_size=1,
shuffle=False)
model = TMPI(num_planes=config.num_planes)
model = model.cuda().eval()
model = torch.nn.parallel.DataParallel(model, device_ids=range(torch.cuda.device_count() ))
model_file = 'weights/mpti_04.pth'
model.load_state_dict( torch.load( os.path.join(os.getcwd(), model_file) ) )
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
✔ 입력 폴더(indir)에서 이미지를 불러와 데이터셋을 만듦
✔ DataLoader를 이용해 한 개씩(batch_size=1) 처리
📍 TMPI 모델 불러오기
✔ TMPI 모델을 GPU에 로드하여 실행
✔ 사전 학습된 mpti_04.pth를 불러옴
📍 이미지 3D 변환 실행
for batch, sample in enumerate(dataset_loader):
if batch < 1:
continue
with torch.no_grad():
poses_tgt = render_path()
pred_frames = render_3d_photo(model, sample, poses_tgt, use_gl_renderer)
pred_frames = [ np.clip(np.round(frame * 255), a_min=0, a_max=255).astype(np.uint8) for frame in pred_frames ]
if crop_edges:
pred_frames = [ frame[20:-20, 20:-20, :] for frame in pred_frames ]
rgb_clip = ImageSequenceClip(pred_frames, fps=30)
rgb_clip.write_videofile( os.path.join(outdir, '%s.mp4' % sample['filename'][0]), verbose=False, codec='mpeg4', logger=None, bitrate='2000k')
✔ render_path()를 이용해 새로운 시점 생성
✔ render_3d_photo()를 이용해 각 시점에서 3D 이미지 생성
✔ 여러 장의 프레임을 만들고 pred_frames에 저장
📍 영상(mp4) 파일 저장
✔ 프레임을 정리하여 MP4 형식으로 저장
✔ fps=30 → 초당 30 프레임으로 자연스럽게 3D 영상 생성
✔ output/ 폴더에 새로운 시점에서의 3D 효과가 포함된 영상 저장
TMPI/dataset.py
TMPI 시스템에서 사용할 이미지 데이터를 불러오고, 이를 3D 변환에 적합한 형태로 전처리
✅ 이미지 데이터를 로드 및 크기 조정
✅ DPT를 이용하여 깊이 정보(Depth Map) 생성
✅ 이미지를 작은 타일(Tile)로 분할
✅ TMPI에서 사용할 수 있도록 정리하여 반환
📌 코드 개요
• TMPIDataset 클래스는 이미지 데이터셋을 관리하고, 깊이 정보를 추출하여 타일(tile) 단위로 분할하는 역할을 합니다.
• DPTWrapper를 이용해 단일 이미지에서 깊이 정보(Depth Map)를 예측합니다.
• OpenGL 기반의 TMPI 시스템에서 사용할 수 있도록 데이터를 변환합니다.
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import os
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
import numpy as np
import glob
import config
from dpt_wrapper import DPTWrapper
from utils import imutils
from utils import utils
import config
✔ DPTWrapper → DPT(Depth Prediction Transformer)를 통해 깊이 정보 추출
✔ imutils → 이미지 변환 유틸리티
✔ utils → 보조 함수 모음
🔹 2) TMPIDataset 클래스 초기화 (__init__)
class TMPIDataset(Dataset):
def __init__(self,
data_root,
image_ext=["png", "jpeg", "jpg"],
):
super(TMPIDataset, self).__init__()
self.data_root = data_root
self.image_path_list = []
for ext in image_ext:
self.image_path_list.extend( glob.glob(self.data_root + "/**/*.%s" % ext, recursive=True) )
self.image_path_list.sort()
self.monocular_depth = DPTWrapper(model_path='./DPT/weights/dpt_hybrid-midas-501f0c75.pt')
✔ 이미지 데이터 경로 읽기
• data_root에 있는 png, jpeg, jpg 파일을 모두 검색하여 리스트에 저장
• glob.glob(..., recursive=True)을 사용하여 하위 폴더까지 탐색
• .sort()를 사용하여 정렬된 리스트 생성
✔ DPT(Depth Prediction Transformer) 모델 불러오기
• DPTWrapper를 사용하여 단일 이미지에서 깊이 정보(Depth Map)를 예측할 수 있도록 설정
• dpt_hybrid-midas-501f0c75.pt라는 사전 학습된 모델을 로드하여 깊이 예측
🔹 3) 타일링(tiles) - 이미지를 겹치는 작은 조각으로 나누기
입력 이미지를 겹치는 작은 타일(tile)로 분할
# Subdivide the image into over-lapping tiles
def tiles(self, src_disp, src_rgb, K, tile_sz, pad_sz):
bs, _, h, w = src_disp.shape
K_, sx_, sy_, dmap_, rgb_ = [], [], [], [], []
sy = torch.arange(0, h, tile_sz - pad_sz)
sx = torch.arange(0, w, tile_sz - pad_sz)
src_disp = F.pad(src_disp, (0, tile_sz, 0, tile_sz), 'replicate')
src_rgb = F.pad(src_rgb, (0, tile_sz, 0, tile_sz), 'replicate')
K_, src_disp_, src_rgb_, sx_, sy_ = [], [], [], [], []
✔ bs, h, w → 입력 이미지의 배치(batch), 높이(height), 너비(width)
✔ sy, sx → 이미지에서 타일의 시작 좌표
✔ F.pad(..., 'replicate') → 이미지 테두리를 복제하여 패딩 추가
✔ 패딩을 추가하는 이유는 타일을 만들 때 가장자리 부분을 보정하기 위해
✔ F.pad()는 이미지 전체에 한 번만 추가됨.
✔ 각 타일에 개별적으로 패딩이 적용되는 것이 아니라, 패딩이 추가된 전체 이미지에서 타일을 슬라이딩 윈도우 방식으로 추출하는 과정에서 타일이 결정됨.
✔ 이 과정이 없으면, 경계 부분의 타일이 원본 이미지보다 작은 크기로 잘릴 수 있음.
for y in sy:
for x in sx:
l, r, t, b = x, x + tile_sz, y, y + tile_sz
Ki = K.clone()
Ki[:, 0, 2] = Ki[:, 0, 2] - x
Ki[:, 1, 2] = Ki[:, 1, 2] - y
K_.append(Ki)
src_disp_.append( src_disp[:, :, t:b, l:r] )
src_rgb_.append( src_rgb[:, :, t:b, l:r] )
sx_.append(x)
sy_.append(y)
src_rgb_ = torch.stack(src_rgb_, 1)
src_disp_ = torch.stack(src_disp_, 1)
K_ = torch.stack(K_, 1)
sx_, sy_ = torch.tensor(sx_).unsqueeze(0).expand(bs, -1), torch.tensor(sy_).unsqueeze(0).expand(bs, -1)
return src_disp_, src_rgb_, K_, sx_, sy_
✔ 각 타일을 선택하고 카메라 행렬(K)도 보정
✔ Ki[:, 0, 2] -= x → X 방향으로 이동한 만큼 원근 투영을 보정
✔ Ki[:, 1, 2] -= y → Y 방향으로 이동한 만큼 보정
✔ 여러 개의 타일을 하나의 텐서(tensor)로 변환하여 저장
✔ src_rgb_ → RGB 이미지의 타일
✔ src_disp_ → 깊이 정보(Depth Map)의 타일
🔹 4) 데이터셋 크기 반환 (__len__)
def __len__(self):
return len(self.image_path_list)
✔ len(dataset)을 호출하면 데이터셋에 포함된 이미지 개수를 반환
🔹 5) 데이터셋에서 이미지 로드 (__getitem__)
def __getitem__(self, idx):
src_rgb = imutils.png2np( self.image_path_list[idx] ).astype(np.float32)
h, w = src_rgb.shape[:2]
# Scale the image if too large
if h >= w and h >= config.imgsz_max:
h_scaled, w_scaled = config.imgsz_max, int(config.imgsz_max / h * w)
elif w > h and w >= config.imgsz_max:
h_scaled, w_scaled = int(config.imgsz_max / w * h), config.imgsz_max
else:
h_scaled, w_scaled = h, w
✔ self.image_path_list[idx] → 해당 인덱스의 이미지 파일을 로드
✔ imutils.png2np(...) → 이미지를 numpy 배열로 변환
📍 이미지 크기 조정 (과도하게 크면 줄이기)
✔ config.imgsz_max → 이미지 크기가 너무 크면 자동으로 크기를 줄임
🔹 6) 깊이 정보 추출
src_rgb = F.interpolate(torch.from_numpy(src_rgb).permute(2, 0, 1).unsqueeze(0), (h_scaled, w_scaled), mode="bilinear")
# Estimate monocular depth
src_disp = torch.from_numpy(self.monocular_depth( src_rgb.squeeze().permute(1, 2, 0).numpy())).unsqueeze(0).unsqueeze(0)
src_disp = (src_disp - torch.min(src_disp)) / (torch.max(src_disp) - torch.min(src_disp)) # Normalize to [0, 1]✔ DPTWrapper를 사용하여 깊이 정보(Depth Map) 생성
✔ 깊이 정보를 [0,1] 범위로 정규화(Normalization)
🔹 7) 카메라 행렬(K) 설정
K = torch.tensor([
[0.58, 0, 0.5],
[0, 0.58, 0.5],
[0, 0, 1]
]).unsqueeze(0)
K[:, 0, :] *= w_scaled
K[:, 1, :] *= h_scaled✔ K는 카메라 내부 파라미터 행렬(Intrinsic Matrix)
✔ 이미지 크기에 맞게 K 값을 조정
🔹 8) 타일링 수행 및 최종 반환
tile_sz = int(np.clip(utils.next_power_of_two(config.tilesz2w_ratio * w_scaled - 1), a_min=config.tilesz_min, a_max=config.tilesz_max))
pad_sz = int(tile_sz * config.padsz2tile_ratio)
src_disp_tiles, src_rgb_tiles, K_tiles, sx, sy = self.tiles(src_disp, src_rgb, K, tile_sz, pad_sz)
return {
"src_rgb_tiles": src_rgb_tiles.squeeze(0),
"src_disp_tiles": src_disp_tiles.squeeze(0),
"K": K_tiles.squeeze(0),
"sx": sx.squeeze(0),
"sy": sy.squeeze(0),
"src_rgb": src_rgb.squeeze(0),
"src_disp": src_disp.squeeze(0),
"src_width": w,
"src_height": h,
"tile_sz": tile_sz,
"pad_sz": pad_sz,
"cam_int": K.squeeze(0),
"cam_ext": torch.eye(4),
"filename": os.path.basename(os.path.splitext(self.image_path_list[idx])[0])
}
✔ 타일링한 데이터와 관련 정보들을 딕셔너리 형태로 반환
TMPI/kmeans.py
import torch
import numpy as np
class KMeans:
""" Weighted K-Means clustering """
def __init__(self, k=4, featsz=1, batchsz=1, device='cpu'):
super().__init__() # 부모 클래스 초기화 (필수는 아님)
self.k = k
self.device = device
# Start off with a linear distribution for the cluster centers
self.cluster_m = torch.linspace(0, 1, k).view(1, k, 1).expand(batchsz, -1, featsz).to(device)
def train( self, samples, conf ):
b, n, m = samples.shape
o = torch.randperm(n) # 0부터 n-1까지의 랜덤 순열 생성 (데이터 순서 섞기)
samples = samples[:, o, :] # 샘플을 무작위로 섞기 (Shuffle)
conf = conf[:, o, :] # 신뢰도(confidence)도 동일하게 섞기
for j in range(50): # K-Means 반복 학습 (최대 50번 수행)
# Calculate distance from each sample to each cluster center
samples_ = samples.unsqueeze(-2).expand(-1, -1, self.k, -1) # (b, n, k, m) # 샘플 차원 확장
cluster_m_ = self.cluster_m.unsqueeze(-3).expand(-1, n, -1, -1) # (b, n, k, m) # 클러스터 중심 차원 확장
d = torch.square(samples_ - cluster_m_).sum(-1) # 유클리드 거리 계산 (L2 distance)
w_idx = torch.min(d, dim=-1)[1].view(b, n, 1).expand(-1, -1, m)
# Calculate new cluster centers as the weighted mean of the closest samples
for i in range(self.k): # 각 클러스터에 대해 업데이트 수행
w_i = w_idx == i # 해당 클러스터에 속한 샘플을 선택
self.cluster_m[:, i, :] = torch.sum(samples * conf * w_i, 1) / torch.sum(w_i * conf, 1).clamp(min=1)• k: 클러스터 개수 (기본값: 4)
• featsz: 특징 차원 수 (default: 1)
• batchsz: 배치 크기 (default: 1)
✅ 초기 클러스터 중심(cluster_m)을 선형 분포로 설정
1. torch.linspace(0, 1, k): 0에서 1까지 k개의 등간격 값을 생성 (즉, 클러스터 중심 초기화),
생성된 값들은 torch.Tensor 형태로 반환
2. .view(1, k, 1): (1, k, 1) 크기로 변환 → 배치 크기 적용을 위한 차원 확장
- • 텐서의 차원(Shape)을 변경
- • torch.linspace(0, 1, k)의 출력은 1D 텐서 (k,)
- • .view(1, k, 1)을 적용하면 (1, k, 1)로 변경 (배치 크기와 feature 차원 추가)
3. .expand(batchsz, -1, featsz): batchsz 개의 클러스터를 동일한 초기값으로 설정,
-1은 기존 크기를 유지한다는 의미
4. .to(device): GPU 또는 CPU에서 연산할 수 있도록 설정
✅ 클러스터링 학습 (Weighted K-Means Training)
- • samples: 입력 샘플 데이터 (batch_size, num_samples, feature_dim)
- • samples는 K-Means 클러스터링에 입력되는 데이터
- • batch_size: 한 번에 처리할 데이터 그룹 개수
- • num_samples: 각 배치 내 샘플 개수
- • feature_dim: 각 샘플이 가지는 특징 차원 (예: RGB라면 3, 깊이면 1)
- • conf: 각 샘플의 신뢰도 (batch_size, num_samples, feature_dim)
✅ 입력 데이터의 차원 정보 저장
• b: 배치 크기
• n: 샘플 개수
• m: 특징 차원
✅ 샘플 순서 랜덤화 (Shuffling)
- • K-Means는 초기 클러스터 배치에 따라 결과가 달라질 수 있음
- • 같은 데이터라도 중심이 어디서 시작되냐에 따라 다른 군집 결과를 낼 수 있음
- • 따라서 초기 데이터 순서를 무작위로 섞어(K-Means++) 최적의 클러스터 중심을 찾도록 함
✅ K-Means 알고리즘의 반복 횟수 설정
• 일반적인 K-Means처럼 중심이 수렴할 때까지 반복
• j는 반복 횟수를 의미하며, 최대 50번 반복
✅ 각 샘플과 클러스터 중심 간의 거리 계산
1. samples.unsqueeze(-2): 샘플 차원을 확장해 (b, n, 1, m) 형태로 변경
- Tensor.unsqueeze(dim) -> Tensor
- • dim: 추가할 차원의 위치
- • -1: 마지막 차원 추가
- • -2: 끝에서 두 번째 차원 추가
- • -3: 끝에서 세 번째 차원 추가
2. .expand(-1, -1, self.k, -1): k개의 클러스터와 비교할 수 있도록 확장 (b, n, k, m)
3. cluster_m.unsqueeze(-3): 클러스터 중심도 동일한 차원 (b, n, k, m)으로 확장
4. torch.square(samples_ - cluster_m_).sum(-1): L2 거리 (유클리드 거리의 제곱) 계산
✅ 각 샘플을 가장 가까운 클러스터에 할당
• dim=-1은 마지막 차원에서 연산을 수행한다는 의미
• sum(-1): 마지막 차원의 값을 모두 더함
• torch.min(d, dim=-1)[1]: 거리 d에서 가장 작은 값을 가지는 클러스터 인덱스를 찾음
• .view(b, n, 1): 차원 조정 후
• .expand(-1, -1, m): feature dimension을 동일하게 확장
✅ 새로운 클러스터 중심 업데이트 (Weighted Mean)
• w_i = w_idx == i: 클러스터 i에 속한 샘플을 마스크
- • w_i는 현재 클러스터 i에 속한 데이터만 선택하는 역할
- • True/False의 Boolean Mask로 동작
- • 마스킹된 값만 남기고 나머지는 제거
- 예시
labels = torch.tensor([0, 1, 0, 2, 1])
w_i = labels == 1 # 클러스터 1에 속한 데이터 찾기
print(w_i) # tensor([False, True, False, False, True])• samples * conf * w_i: 해당 클러스터의 가중치를 적용한 샘플
• torch.sum(samples * conf * w_i, 1) / torch.sum(w_i * conf, 1).clamp(min=1):
• 각 클러스터의 새로운 중심을 **가중 평균(Weighted Mean)**으로 계산
• .clamp(min=1): 0으로 나누는 문제 방지
• Weighted K-Means에서는 각 샘플에 대한 Confidence 값(신뢰도)를 가중치로 적용
• conf를 곱하지 않으면 일반적인 K-Means Clustering이 됨
def test( self, samples ):
b, n, m = samples.shape # 배치 크기, 샘플 개수, 특징 차원
# 샘플과 클러스터 중심을 비교할 수 있도록 차원 확장
samples = samples.unsqueeze(-2).expand(-1, -1, self.k, -1)
cluster_m = self.cluster_m.unsqueeze(-3).expand(-1, n, -1, -1)
#각 샘플을 가장 가까운 클러스터로 매칭
d = torch.square(samples - cluster_m).sum(-1) # 거리 계산
w_idx = torch.min(d, dim=-1)[1].view(b, n, 1).expand(-1, -1, m) # 가장 가까운 클러스터 찾기
w = torch.gather(self.cluster_m, 1, w_idx) # 해당 클러스터의 중심값을 가져오기
# Return the cluster label, and cluster index for each sample
return w.view(b, n, m), w_idx[:, :, 0] # (샘플의 클러스터 중심값, 클러스터 인덱스)
클러스터링 테스트 (Inference)
• 새로운 입력 데이터에 대해 클러스터 예측 수행
✅ 클러스터링 결과 반환
• w: 각 샘플이 속한 클러스터의 중심값
• w_idx: 클러스터 인덱스
<참고>
• 단순 평균은 모든 샘플을 동일한 가중치로 반영하지만,
• 가중 평균은 신뢰도가 높은 샘플을 더 반영하여 클러스터 중심을 더 정확하게 계산
# 일반 K-Means 평균
new_center = torch.sum(samples * w_i, 1) / torch.sum(w_i, 1)
# Weighted K-Means (가중 평균)
new_center = torch.sum(samples * conf * w_i, 1) / torch.sum(w_i * conf, 1).clamp(min=1)• samples * w_i: 특정 클러스터에 속한 샘플들만 선택 (입력데이터* 각 샘플이 속한 클러스터의 인덱스의 마크스(t/f))
- 각 샘플에 대해 w_i가 1인 샘플들만 남기고, 나머지는 0으로 만듦
- 즉, 현재 계산 중인 클러스터(i)에 속한 샘플들만 남김
• torch.sum(w_i, 1): 해당 클러스터에 속한 샘플 개수를 계산
• 즉, w_i가 1(True)인 샘플들의 평균을 구하는 코드
✅ 왜 clamp(min=1)이 필요 없나?
• 일반적인 K-Means에서는 분모가 0이 되는 경우가 거의 없음.
• K-Means는 각 클러스터에 최소한 하나의 샘플이 할당되도록 동작
• 따라서, torch.sum(w_i, 1)이 0이 될 가능성이 거의 없음
• K-Means++ 초기화를 사용할 경우, 클러스터가 고르게 분배되도록 설정됨 → 분모가 0이 될 가능성이 낮음
TMPI/tmpi.py
• 깊이 신뢰도 예측 (self.conf_net) → 깊이 및 신뢰도 맵 생성
• 타일링 (F.unfold) → 입력 데이터를 타일 단위로 변환
• K-Means 클러스터링 (KMeans) → 타일 내에서 깊이 평면을 구분
• 정규화 및 복원 (depth_tiles_norm) → 깊이를 0~1 범위로 정규화 후 복원
• MPI 생성 및 반환 → 최종 다층 평면 영상(MPI) 생성
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
# 필요한 라이브러리 import
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import config
import numpy as np
from utils import utils
from kmeans import KMeans # K-Means 클러스터링 사용
import networks # 신경망 모델 포함
# TMPI 클래스 정의 (Tiled MultiPlane Image)
class TMPI(torch.nn.Module):
def __init__(self, num_planes = 4):
""" TMPI 모델 초기화 """
super().__init__()
self.num_planes = num_planes # 다중 평면 개수 설정
# 신뢰도 예측 네트워크 (깊이 신뢰도 예측)
self.conf_net = networks.DepthConfidenceNet()
# MPI 네트워크 (MPI 생성)
self.mpti_net = networks.TileMPINet( cin=num_planes * 4 + 1,
c0=32,
cout=num_planes * 4,
depth=3)
def inpaint(self, rgb, depth, feature_masks):
""" 심도 및 RGB 데이터를 기반으로 보간하여 손실된 영역을 복원하는 함수 """
b, _, h, w = rgb.shape # 배치 크기, 높이, 너비 추출
# feature_masks: 한 평면에서 가려진 영역을 다음 평면에서 마스킹
feature_masks = torch.cat( (torch.ones_like(feature_masks[:, 0:1, ...]).to(feature_masks.get_device()),
feature_masks[:, :-1, ...]), 1)
feature_masks = 1 - feature_masks # 알파 마스크를 가시성 마스크로 변환
# 경계 영역 마스킹 (깊이 변화가 적은 부분만 보존)
feature_masks = feature_masks * (utils.gradient(depth) < 0.01).float().expand(-1, feature_masks.shape[1], -1, -1)
# 첫 번째 평면은 항상 가시성 1
feature_masks[:, 0, ...] = 1
feature_masks = feature_masks.view(b, config.num_planes - 1, 1, h, w).expand(-1, -1, 3, -1, -1)
# 원본 RGB 데이터에 가시성 마스크 적용
rgb_m = rgb.view(b, 1, 3, h, w).expand(-1, config.num_planes - 1, -1, -1, -1) * feature_masks
# 평균 풀링 레이어 정의
sumpool2d = torch.nn.AvgPool2d(kernel_size=2, stride=2, padding=0, divisor_override=1)
# 재귀적 보간을 수행하는 함수 (피라미드 구조 활용)
def inpaint_level(rgb_, fmask, level):
if level == 0 or rgb_.shape[-1] == 1 or rgb.shape[-2] == 1:
return rgb_
return rgb_ * fmask + (1 - fmask) * F.interpolate(
inpaint_level(sumpool2d(rgb_) / sumpool2d(fmask).clamp(min=1),
sumpool2d(fmask).clamp(max=1),
level - 1),
size = rgb_.shape[-2:], mode='bilinear')
# 최종 보간된 결과 반환
return inpaint_level(rgb_m.view(b, -1, h, w), feature_masks.reshape(b, -1, h, w), np.log2(h))forward 함수는 모델이 입력을 받아 연산을 수행하는 주요 부분이다.
• 입력값:
• rgb_tiles: (B, N, 3, H, W) → RGB 이미지 타일
• inv_depth_tiles: (B, N, 1, H, W) → 깊이 정보가 포함된 타일 (역깊이)
• rgb: 전체 RGB 이미지 (B, 3, H, W)
• inv_depth: 전체 역깊이 이미지 (B, 1, H, W)
• tile_sz: 타일의 크기.
• pad_sz: config.padsz2tile_ratio를 사용하여 패딩 크기 계산.
• replicate: 가장자리 픽셀을 복제하여 패딩 수행.
• conf_net: DepthConfidenceNet을 사용하여 깊이(pred_depth) 및 신뢰도(pred_conf) 예측.
• torch.cat((rgb_padded, inv_depth_padded), 1): RGB와 깊이 정보를 결합하여 네트워크에 입력.
• sy, sx: 타일 생성 위치를 나타내는 인덱스 배열.
• F.unfold: pred_depth_conf에서 슬라이딩 윈도우 방식으로 타일을 추출.
• kernel_size=tile_sz: 타일 크기만큼 가져옴.
• stride=(tile_sz - pad_sz): 겹쳐지는 영역을 포함하여 타일을 이동.
• depth_conf_tiles의 차원을 재구성하여 타일 데이터를 (B, 타일 개수, 2, H, W) 형태로 변환.
• KMeans를 사용하여 깊이 타일을 num_planes - 1 개의 클러스터로 분할.
• depth_tiles_norm과 conf_tiles를 (B*N, H*W, 1) 형태로 변환 후 K-Means 학습.
• 학습된 K-Means 모델을 사용하여 각 타일의 픽셀을 클러스터에 할당.
• 클러스터링된 깊이 값을 다시 원래 깊이 범위로 변환.
• dplanes: K-Means의 클러스터 중심값을 기반으로 깊이 평면 생성.
• 클러스터링된 깊이값을 기반으로 평면 마스크 생성.
• plane_masks의 누적 합을 사용하여 가려진 부분을 계산.
def forward(self, rgb_tiles, inv_depth_tiles, rgb, inv_depth): # 모델의 순전파 과정 정의
device = inv_depth.get_device() # 입력된 깊이 데이터가 위치한 디바이스 확인 (CPU 또는 CUDA)
h, w = rgb.shape[-2:] # 전체 RGB 이미지의 높이(h)와 너비(w)를 가져옴
b, n, _, ht, wt = inv_depth_tiles.shape # 타일별 깊이 데이터의 배치 크기(b), 타일 개수(n), 채널(무시), 타일 높이(ht), 타일 너비(wt) 확인
tile_sz = ht # 타일 크기를 할당
pad_sz = int(tile_sz * config.padsz2tile_ratio) # 패딩 크기를 config 값을 기준으로 계산
# Pad images by tile_sz so that we have a full tile for each image pixel
rgb_padded = F.pad(rgb, (0, tile_sz, 0, tile_sz), 'replicate') # RGB 이미지의 가장자리 픽셀을 복제하여 패딩 적용
inv_depth_padded = F.pad(inv_depth, (0, tile_sz, 0, tile_sz), 'replicate') # 깊이 이미지도 동일하게 패딩 적용
# Pad to multiple of 16 (for network computation)
padder = utils.InputPadder(rgb_padded.shape[-2:], divis_by=16) # 네트워크 연산을 위해 크기를 16의 배수로 맞춤
rgb_padded, inv_depth_padded = padder.pad(rgb_padded, inv_depth_padded) # 패딩 적용된 RGB 및 깊이 이미지 반환
pred_depth, pred_conf = self.conf_net(torch.cat((rgb_padded, inv_depth_padded), 1)) # DepthConfidenceNet을 사용하여 깊이(pred_depth)와 신뢰도(pred_conf) 예측
pred_depth = padder.unpad(pred_depth) # 패딩된 영역 제거하여 원본 크기로 복원
pred_conf = padder.unpad(pred_conf) # 패딩된 영역 제거하여 원본 크기로 복원
pred_depth_conf = torch.cat((pred_depth, pred_conf), 1) # 깊이와 신뢰도를 하나의 텐서로 결합
# Generate tiles from predicted depth and confidence
sy = torch.arange(0, h, tile_sz - pad_sz) # 타일 위치를 결정하기 위한 y축 인덱스 배열 생성
sx = torch.arange(0, w, tile_sz - pad_sz) # 타일 위치를 결정하기 위한 x축 인덱스 배열 생성
depth_conf_tiles = F.unfold(pred_depth_conf[:, :, :(sy[-1] + tile_sz), :(sx[-1] + tile_sz)], # 깊이 및 신뢰도를 타일 단위로 변환
kernel_size=tile_sz,
stride=(tile_sz - pad_sz)) # 타일 크기만큼 슬라이딩 윈도우 적용
depth_conf_tiles = depth_conf_tiles.view(b, 2, tile_sz, tile_sz, -1).permute(0, 4, 1, 2, 3) # 차원 변환하여 타일 데이터 정리
depth_tiles = depth_conf_tiles[:, :, 0, None, ...] # 깊이 정보만 분리
conf_tiles = depth_conf_tiles[:, :, 1, None, ...] # 신뢰도 정보만 분리
# Normalize each depth tile independently to [0, 1]
mn_t = torch.min(depth_tiles.view(b, n, -1), -1)[0].view(b, n, 1, 1, 1) # 각 타일의 최소 깊이값 계산
mx_t = torch.max(depth_tiles.view(b, n, -1), -1)[0].view(b, n, 1, 1, 1) # 각 타일의 최대 깊이값 계산
depth_tiles_norm = (depth_tiles - mn_t) / (mx_t - mn_t + 1e-10) # 깊이값을 0~1 범위로 정규화하여 K-Means 학습을 안정화
# K-means clustering for plane depth placement within each tile
classifier = KMeans(k=config.num_planes - 1, featsz=1, batchsz=b * n, device=device) # K-Means 클러스터링 모델 초기화
classifier.train(depth_tiles_norm.permute(0, 1, 3, 4, 2).reshape(b * n, ht * wt, -1), # K-Means 학습 수행
conf_tiles.permute(0, 1, 3, 4, 2).reshape(b * n, ht * wt, -1)) # 신뢰도와 함께 학습 수행
labels, cluster_idx = classifier.test(depth_tiles_norm.permute(0, 1, 3, 4, 2).reshape(b * n, ht * wt, -1)) # 클러스터링된 결과 얻기
# Use each tile pixel's label as the assigned plane depth, and unnormalize to original tile depth range
depth_assign_tiles = labels.permute(0, 2, 1).unsqueeze(-1) # K-Means 레이블을 깊이 값으로 변환
depth_assign_tiles = depth_assign_tiles.view(b, n, 1, ht, wt) # (B, N, 1, H, W) 형태로 변환
depth_assign_tiles = depth_assign_tiles * (mx_t - mn_t) + mn_t # 정규화된 값을 원래 깊이 범위로 복원
cluster_idx = cluster_idx.view(b, n, 1, ht, wt).expand(-1, -1, config.num_planes - 1, -1, -1) # 클러스터 인덱스를 평면 개수에 맞게 확장
dplanes = classifier.cluster_m.view(b, n, config.num_planes - 1).flip(dims=[-1]).sort(-1, descending=True)[0] # K-Means 클러스터 중심을 기반으로 깊이 평면을 생성
dplanes = dplanes * (mx_t.view(b, n, 1) - mn_t.view(b, n, 1)) + mn_t.view(b, n, 1) # 정규화된 깊이 값을 원래 범위로 변환
plane_masks = cluster_idx - torch.arange(0, config.num_planes - 1).flip(-1).to(device).view(1, 1, config.num_planes - 1, 1, 1).expand(b, n, -1, -1, -1) # 평면 마스크 생성
plane_masks = (plane_masks == 0).float() # 마스크를 이진 형태로 변환
occlusion_masks = torch.cumsum(plane_masks, dim=2).clamp(max=1) # 누적 마스크를 사용하여 가려진 부분을 계산
# Get a prior for what the MPI tiles should look like
mpi0_rgb_tiles = self.inpaint(rgb_tiles.view(b * n, 3, ht, wt), # RGB 이미지 타일에 대한 인페인팅 수행
inv_depth_tiles.view(b * n, 1, ht, wt), # 깊이 이미지 타일을 활용
occlusion_masks.clone().detach().view(b * n, config.num_planes - 1, ht, wt)).clamp(0, 1) # 가려진 부분을 보완
mpi0_rgba_tiles = torch.cat((mpi0_rgb_tiles.view(b, n, config.num_planes - 1, -1, ht, wt),
occlusion_masks.view(b, n, -1, 1, ht, wt)), 3) # RGB + Alpha 채널 결합
# 마지막 레이어에 투명도 채널 추가
mpi0_rgba_tiles = torch.cat((mpi0_rgba_tiles, mpi0_rgba_tiles[:, :, -1, None, ...]), 2)
mpi0_rgba_tiles[:, :, -1, -1, ...] = 1.0 # 마지막 레이어를 완전히 불투명하게 설정
dplanes = torch.cat((dplanes, torch.ones_like(dplanes[:, :, -1, None]).to(device) * 1e-7), -1) # 깊이 평면에 추가 레이어 적용
padder = utils.InputPadder((ht, wt), divis_by=16) # 깊이 타일과 MPI 데이터를 패딩하여 16의 배수로 변환
inv_depth_tiles, mpi0_rgba_tiles = padder.pad(inv_depth_tiles.view(b * n, -1, ht, wt),
mpi0_rgba_tiles.view(b * n, config.num_planes * 4, ht, wt)) # 패딩 적용
mpi_tiles = F.sigmoid(self.mpti_net(torch.cat((mpi0_rgba_tiles, inv_depth_tiles), 1))) # MPI 타일을 최종 네트워크에 입력하여 결과 생성
mpi_tiles = padder.unpad(mpi_tiles).view(b, n, self.num_planes, 4, ht, wt).contiguous() # 패딩 제거 후 최종 형태로 변환
# Blend prediction with prior based on alpha
blendw_tiles = torch.cumprod(1 - mpi_tiles[:, :, :, -1, None, ...], 2) # 알파 채널을 기반으로 블렌딩 가중치 생성
blendw_tiles = torch.cat([torch.ones_like(blendw_tiles[:, :, 0, None, ...]).cuda(), blendw_tiles[:, :, :-1, ...]], 2) # 첫 번째 채널을 유지
return mpi_tiles, dplanes # 최종적으로 MPI 타일과 깊이 평면 반환
1. 정규화 관련 코드
mn_t = torch.min(depth_tiles.view(b, n, -1), -1)[0].view(b, n, 1, 1, 1)
mx_t = torch.max(depth_tiles.view(b, n, -1), -1)[0].view(b, n, 1, 1, 1)
의미:
• mn_t: depth_tiles의 최소값을 각 타일(tile)별로 계산한 결과.
• depth_tiles.view(b, n, -1)은 각 타일의 픽셀 값을 (배치 크기, 타일 개수, 픽셀 값) 형태로 펼친 후,
• torch.min(..., -1)[0]은 마지막 차원(픽셀 값)에서 최소값을 계산합니다.
• .view(b, n, 1, 1, 1)은 다시 원래 차원으로 복원합니다.
• mx_t: depth_tiles의 최대값을 각 타일(tile)별로 계산한 결과. 계산 과정은 mn_t와 동일하지만 최대값(torch.max)을 계산합니다.
목적: 각 타일별 깊이 값의 범위를 계산해, 이를 기반으로 정규화를 수행하거나 복원을 진행합니다.
2. 정규화 되돌리기 (Unnormalization)
depth_assign_tiles = depth_assign_tiles * (mx_t - mn_t) + mn_t
의미:
• depth_assign_tiles는 정규화된 상태(0~1 범위)로 나타난 데이터를 원래의 절대값 범위로 되돌립니다.
• (mx_t - mn_t)는 해당 타일의 원래 깊이 범위 (최대값 - 최소값)이고,
• + mn_t는 타일의 최소값을 더해 원래 깊이 값으로 복원합니다.
목적: 정규화된 깊이 데이터를 원래의 절대 깊이 값으로 변환하여, 실제 깊이와 동일한 범위로 만들기 위함.
3. dplanes 복원
dplanes = dplanes * (mx_t.view(b, n, 1) - mn_t.view(b, n, 1)) + mn_t.view(b, n, 1)
의미:
• dplanes는 K-means 클러스터링 결과로 나온 타일의 평균 깊이(plane depth)를 나타냅니다.
• 이를 정규화 상태에서 원래 깊이 범위로 복원합니다:
• (mx_t.view(b, n, 1) - mn_t.view(b, n, 1))는 깊이 값의 범위,
• + mn_t.view(b, n, 1)는 최소값을 더해 원래 깊이 범위로 변환합니다.
목적: 클러스터링에서 사용된 dplanes 값을 실제 깊이 값으로 변환하여 절대적인 깊이 정보를 재구성합니다.
mpi_tiles.shape = (b, n, num_planes, 4, ht, wt)
@@@ mpi_tiles : torch.Size([1, 54, 4, 4, 64, 64])