jens/pyRBM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[faces_sub_image] - add --grayscale CLI flag for single-channel mode

Browse Source 
GRAYSCALE flag drives N_CH (1 or 3) and N_VIS (1024 or 3072) at runtime.
_img_to_chw() handles BGR→gray or BGR→RGB conversion.
All visualisation functions use cmap='gray' when GRAYSCALE is set.
Separate model name (faces_sub_image_gray) keeps weights isolated.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jens
2026-05-30 16:07:08 +02:00
co-authored by Claude Sonnet 4.6
parent fc10fab627
commit 07fc7dc607
1 changed files with 40 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/tests/test_faces_sub_image.py
+40 -21
View File
@@ -8,13 +8,14 @@ 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
STRIDE = 16 # 50% overlap; set equal to PATCH for non-overlapping
N_CH = 3
N_VIS = N_CH * PATCH * PATCH # 3072
N_HID = 128
N_IMAGES = 50
DATA_DIR = '/media/jens/cifs/bilder/MachineVision/Caltech_WebFaces/'
PATCH = 32
STRIDE = 16 # 50% overlap; set equal to PATCH for non-overlapping
GRAYSCALE = False # reassigned in __main__ when --grayscale is set
N_CH = 1 if GRAYSCALE else 3
N_VIS = N_CH * PATCH * PATCH # 1024 grayscale / 3072 colour
N_HID = 128
N_IMAGES = 50
class TestModel(Model):
@@ -50,8 +51,17 @@ def extract_patches_numpy(img_chw: numpy.ndarray) -> numpy.ndarray:
return numpy.stack(rows, axis=0)
def _img_to_chw(img: numpy.ndarray) -> numpy.ndarray:
"""Convert cv2 BGR image to (C, H, W) float64 array, respecting GRAYSCALE flag."""
if GRAYSCALE:
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(numpy.float64) / 255.0
return gray[numpy.newaxis] # (1, H, W)
img_f = img[:, :, ::-1].astype(numpy.float64) / 255.0 # BGR→RGB
return numpy.ascontiguousarray(numpy.transpose(img_f, (2, 0, 1))) # (3, H, W)
def load_patches(data_dir: str, n_images: int = N_IMAGES, start: int = 0):
"""Load RGB images, extract overlapping patches, return normalised cupy array."""
"""Load images, extract overlapping patches, return normalised cupy array."""
all_patches = []
files = sorted(f for f in os.listdir(data_dir) if f.endswith('.jpg'))[start:start + n_images]
for fname in files:
@@ -59,8 +69,7 @@ def load_patches(data_dir: str, n_images: int = N_IMAGES, start: int = 0):
if img is None or img.shape[0] < PATCH or img.shape[1] < PATCH:
continue
img = pad_to_stride(img)
img_f = img[:, :, ::-1].astype(numpy.float64) / 255.0
img_chw = numpy.ascontiguousarray(numpy.transpose(img_f, (2, 0, 1)))
img_chw = _img_to_chw(img)
all_patches.append(extract_patches_numpy(img_chw))
patches_np = numpy.concatenate(all_patches, axis=0)
@@ -78,8 +87,7 @@ def load_image_patches(path: str):
nx_steps = (h - PATCH) // STRIDE + 1
ny_steps = (w - PATCH) // STRIDE + 1
img_f = img[:, :, ::-1].astype(numpy.float64) / 255.0
img_nchw = np.array(numpy.ascontiguousarray(numpy.transpose(img_f, (2, 0, 1))[numpy.newaxis]))
img_nchw = np.array(_img_to_chw(img)[numpy.newaxis])
patches_nchw = sub(img_nchw)
patches = patches_nchw.reshape(-1, N_VIS)
@@ -103,7 +111,7 @@ def assemble_overlapping(patch_list, nx_steps: int, ny_steps: int) -> numpy.ndar
accum[:, x:x+PATCH, y:y+PATCH] += arr
count[0, x:x+PATCH, y:y+PATCH] += 1.0
img = (accum / numpy.maximum(count, 1)).transpose(1, 2, 0) # (H, W, C)
img = (accum / numpy.maximum(count, 1)).transpose(1, 2, 0).squeeze() # (H,W,C) or (H,W)
lo, hi = img.min(), img.max()
return numpy.clip((img - lo) / (hi - lo + 1e-8), 0, 1)
@@ -117,9 +125,10 @@ def show_filters(entity: Entity, rows: int = 8, cols: int = 16):
ax.axis('off')
if j >= N_HID:
continue
filt = W[:, j].reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0)
filt = W[:, j].reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0).squeeze()
lo, hi = filt.min(), filt.max()
ax.imshow(numpy.clip((filt - lo) / (hi - lo + 1e-8), 0, 1))
ax.imshow(numpy.clip((filt - lo) / (hi - lo + 1e-8), 0, 1),
cmap='gray' if GRAYSCALE else None)
plt.tight_layout()
@@ -130,16 +139,18 @@ def show_reconstructions(model: TestModel, patches: Mat, n_show: int = 10):
fig.suptitle('Patch reconstructions (normalised display)', fontsize=10)
err_total = 0.0
cmap = 'gray' if GRAYSCALE else None
def to_img(p):
a = convert(p).reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0)
a = convert(p).reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0).squeeze()
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')
axes[0, i].imshow(to_img(inp), cmap=cmap); axes[0, i].axis('off')
axes[1, i].imshow(to_img(recon), cmap=cmap); axes[1, i].axis('off')
print(f'Mean patch reconstruction MAE: {err_total / n_show:.4f}')
plt.tight_layout()
@@ -152,8 +163,9 @@ def show_image_reconstruction(model: TestModel, img_path: str):
recon_grid = assemble_overlapping(recons, nx_steps, ny_steps)
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
axes[0].imshow(orig_grid); axes[0].set_title('Original (normalised)'); axes[0].axis('off')
axes[1].imshow(recon_grid); axes[1].set_title('Reconstructed (overlap blended)'); axes[1].axis('off')
cmap = 'gray' if GRAYSCALE else None
axes[0].imshow(orig_grid, cmap=cmap); axes[0].set_title('Original (normalised)'); axes[0].axis('off')
axes[1].imshow(recon_grid, cmap=cmap); axes[1].set_title('Reconstructed (overlap blended)'); axes[1].axis('off')
fig.suptitle(os.path.basename(img_path), fontsize=9)
plt.tight_layout()
@@ -165,9 +177,16 @@ if __name__ == '__main__':
help='Load saved model weights (default: true)')
ap.add_argument('--do_train', type=lambda s: s.lower() != 'false', default=False,
help='Train the model (default: false)')
ap.add_argument('--grayscale', action='store_true', default=False,
help='Use single-channel grayscale patches (default: false)')
args = ap.parse_args()
prj_name = 'faces_sub_image'
if args.grayscale:
GRAYSCALE = True
N_CH = 1
N_VIS = PATCH * PATCH # 1024
prj_name = 'faces_sub_image_gray' if args.grayscale else 'faces_sub_image'
work_dir = 'results'
model = TestModel(prj_name, work_dir)