[faces_sub_image] - add GB-RBM autoencoder test on Caltech WebFaces using SubImage ROIs
Extracts non-overlapping 32x32 RGB patches from face images using SubImage, trains a GB-RBM (3072 visible, 128 hidden), and visualises learned filters, patch reconstructions, and full image reconstruction. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude Sonnet 4.6
@@ -0,0 +1,177 @@
|
||||
import os
|
||||
import numpy
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from rbm.model import Model
|
||||
from rbm.entity import Entity, EntityParams, TrainingParams
|
||||
from rbm.image import SubImage, normalize
|
||||
from rbm.matrix import Mat, np, convert
|
||||
|
||||
DATA_DIR = '/media/jens/cifs/bilder/MachineVision/Caltech_WebFaces/'
|
||||
PATCH = 32
|
||||
N_CH = 3
|
||||
N_VIS = N_CH * PATCH * PATCH # 3072
|
||||
N_HID = 128
|
||||
N_IMAGES = 50
|
||||
|
||||
|
||||
class TestModel(Model):
|
||||
def __init__(self, name: str, work_dir: str = '.'):
|
||||
super().__init__(name, work_dir)
|
||||
self.unit1 = Entity(
|
||||
(N_VIS, N_HID),
|
||||
EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
|
||||
TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=3000, mini_batch_size=100)
|
||||
)
|
||||
|
||||
def forward(self, x: Mat) -> Mat:
|
||||
return self.unit1.forward(x)
|
||||
|
||||
def reconstruct(self, x: Mat) -> Mat:
|
||||
return self.unit1.reconstruct(x)
|
||||
|
||||
|
||||
def pad_to_multiple(img: numpy.ndarray, m: int = PATCH) -> numpy.ndarray:
|
||||
h, w = img.shape[:2]
|
||||
return numpy.pad(img, ((0, (-h) % m), (0, (-w) % m), (0, 0)), mode='constant')
|
||||
|
||||
|
||||
def extract_patches_numpy(img_chw: numpy.ndarray, patch: int = PATCH) -> numpy.ndarray:
|
||||
"""Extract non-overlapping patches (step=patch) from (C,H,W) array; returns (N,C*P*P)."""
|
||||
C, H, W = img_chw.shape
|
||||
rows = [img_chw[:, x:x+patch, y:y+patch].flatten()
|
||||
for x in range(0, H - patch + 1, patch)
|
||||
for y in range(0, W - patch + 1, patch)]
|
||||
return numpy.stack(rows, axis=0)
|
||||
|
||||
|
||||
def load_patches(data_dir: str, n_images: int = N_IMAGES):
|
||||
"""Load RGB images, extract non-overlapping 32×32 patches, return normalised cupy array."""
|
||||
all_patches = []
|
||||
files = sorted(f for f in os.listdir(data_dir) if f.endswith('.jpg'))[:n_images]
|
||||
for fname in files:
|
||||
img = cv2.imread(os.path.join(data_dir, fname))
|
||||
if img is None or img.shape[0] < PATCH or img.shape[1] < PATCH:
|
||||
continue
|
||||
img = pad_to_multiple(img)
|
||||
img_f = img[:, :, ::-1].astype(numpy.float64) / 255.0 # BGR→RGB, [0,1]
|
||||
img_chw = numpy.ascontiguousarray(numpy.transpose(img_f, (2, 0, 1)))
|
||||
all_patches.append(extract_patches_numpy(img_chw))
|
||||
|
||||
patches_np = numpy.concatenate(all_patches, axis=0)
|
||||
|
||||
# Drop near-constant patches (std ≈ 0 would give NaN after normalize)
|
||||
std = numpy.std(patches_np, axis=1)
|
||||
patches_np = patches_np[std > 0.01]
|
||||
|
||||
patches = np.array(patches_np)
|
||||
return normalize(patches)
|
||||
|
||||
|
||||
def load_image_patches(path: str):
|
||||
"""Extract normalised patches from a single image using SubImage; return (patches, nx, ny)."""
|
||||
sub = SubImage(PATCH, PATCH, PATCH, PATCH)
|
||||
img = cv2.imread(path)
|
||||
img = pad_to_multiple(img)
|
||||
h, w = img.shape[:2]
|
||||
nx, ny = h // PATCH, w // PATCH
|
||||
|
||||
img_f = img[:, :, ::-1].astype(numpy.float64) / 255.0
|
||||
img_nchw = np.array(numpy.ascontiguousarray(numpy.transpose(img_f, (2, 0, 1))[numpy.newaxis]))
|
||||
patches_nchw = sub(img_nchw) # (nx*ny, 3, 32, 32)
|
||||
patches = patches_nchw.reshape(-1, N_VIS)
|
||||
|
||||
mean = np.mean(patches, axis=1, keepdims=True)
|
||||
std = np.std(patches, axis=1, keepdims=True)
|
||||
std = np.where(std < 0.01, np.ones_like(std), std)
|
||||
return (patches - mean) / std, nx, ny
|
||||
|
||||
|
||||
def show_filters(entity: Entity, rows: int = 8, cols: int = 16):
|
||||
W = convert(entity.state.w_hv) # (N_VIS, N_HID) numpy
|
||||
fig, axes = plt.subplots(rows, cols, figsize=(cols * 1.2, rows * 1.2))
|
||||
fig.suptitle(f'Learned GB-RBM filters ({N_HID} hidden units, {PATCH}×{PATCH} RGB patches)',
|
||||
fontsize=10)
|
||||
for j in range(rows * cols):
|
||||
ax = axes[j // cols][j % cols]
|
||||
ax.axis('off')
|
||||
if j >= N_HID:
|
||||
continue
|
||||
filt = W[:, j].reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0)
|
||||
lo, hi = filt.min(), filt.max()
|
||||
ax.imshow(numpy.clip((filt - lo) / (hi - lo + 1e-8), 0, 1))
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def show_reconstructions(model: TestModel, patches: Mat, n_show: int = 10):
|
||||
fig, axes = plt.subplots(2, n_show, figsize=(n_show * 1.5, 3.5))
|
||||
axes[0, 0].set_ylabel('Original')
|
||||
axes[1, 0].set_ylabel('Recon')
|
||||
fig.suptitle('Patch reconstructions (normalised display)', fontsize=10)
|
||||
err_total = 0.0
|
||||
|
||||
def to_img(p):
|
||||
a = convert(p).reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0)
|
||||
return numpy.clip((a - a.min()) / (a.max() - a.min() + 1e-8), 0, 1)
|
||||
|
||||
for i in range(n_show):
|
||||
inp = patches[i]
|
||||
recon = model.reconstruct(model.forward(inp))
|
||||
err_total += float(np.mean(np.abs(inp - recon)))
|
||||
axes[0, i].imshow(to_img(inp)); axes[0, i].axis('off')
|
||||
axes[1, i].imshow(to_img(recon)); axes[1, i].axis('off')
|
||||
print(f'Mean patch reconstruction MAE: {err_total / n_show:.4f}')
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def show_image_reconstruction(model: TestModel, img_path: str):
|
||||
patches, nx, ny = load_image_patches(img_path)
|
||||
recons = [model.reconstruct(model.forward(patches[i])) for i in range(patches.shape[0])]
|
||||
|
||||
def assemble(patch_list):
|
||||
grid = numpy.zeros((nx * PATCH, ny * PATCH, N_CH))
|
||||
for k, p in enumerate(patch_list):
|
||||
x, y = k // ny, k % ny
|
||||
arr = convert(p).reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0)
|
||||
lo, hi = arr.min(), arr.max()
|
||||
grid[x*PATCH:(x+1)*PATCH, y*PATCH:(y+1)*PATCH] = numpy.clip((arr - lo) / (hi - lo + 1e-8), 0, 1)
|
||||
return grid
|
||||
|
||||
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
|
||||
axes[0].imshow(assemble([patches[i] for i in range(len(recons))])); axes[0].set_title('Original (normalised)'); axes[0].axis('off')
|
||||
axes[1].imshow(assemble(recons)); axes[1].set_title('Reconstructed'); axes[1].axis('off')
|
||||
fig.suptitle(os.path.basename(img_path), fontsize=9)
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
prj_name = 'faces_sub_image'
|
||||
work_dir = 'results'
|
||||
|
||||
model = TestModel(prj_name, work_dir)
|
||||
model.init(0.001)
|
||||
model.load()
|
||||
|
||||
print(f'Loading patches from {N_IMAGES} images...')
|
||||
train_patches = load_patches(DATA_DIR, N_IMAGES)
|
||||
print(f'Training on {train_patches.shape[0]} patches ({N_VIS}-dim each)')
|
||||
|
||||
model.train(train_patches)
|
||||
model.save()
|
||||
|
||||
# Show learned weight filters
|
||||
show_filters(model.unit1)
|
||||
|
||||
# Show patch-level reconstructions on held-out images
|
||||
test_patches = load_patches(DATA_DIR, n_images=N_IMAGES + 10)
|
||||
test_patches = test_patches[train_patches.shape[0]:] # use only unseen patches
|
||||
if test_patches.shape[0] < 10:
|
||||
test_patches = train_patches
|
||||
show_reconstructions(model, test_patches)
|
||||
|
||||
# Show full image reconstruction
|
||||
test_img = os.path.join(DATA_DIR, sorted(os.listdir(DATA_DIR))[60])
|
||||
show_image_reconstruction(model, test_img)
|
||||
|
||||
plt.show()
|
||||
Reference in New Issue
Block a user