diff --git a/src/tests/test_faces_sub_image.py b/src/tests/test_faces_sub_image.py index e022734..c7c8a3e 100644 --- a/src/tests/test_faces_sub_image.py +++ b/src/tests/test_faces_sub_image.py @@ -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)