Added better visualization

app.py

@@ -396,7 +396,7 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
         plt.show()  
 
         return fig 
-    def predict_image(upload,show_original,max_channels):
+    def predict_image(upload,show_original,max_channels,cnn_couple_channel):
         img = cv2.cvtColor(cv2.imread(upload),cv2.COLOR_BGR2RGB)
         model_base_path = "./trained_model"
         classic_model_path =os.path.join(model_base_path,'classic_model.pt')
@@ -412,7 +412,7 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
             return "No CNN Model trained",None,None,None
         cnn_predict = cnn_model.predict(img)
         classic_predict = classic_model.predict(img)
-        cnn_features = cnn_model.visualize_feature(img,max_channels=max_channels)
+        cnn_features = cnn_model.visualize_feature(img,max_channels=max_channels,couple=cnn_couple_channel)
         classical_features = classic_model.visualize_feature(img,show_original=show_original)
         return None,make_figure_from_image(img),cnn_predict,classic_predict,cnn_features,classical_features
     
@@ -423,6 +423,7 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
             gr.Markdown("# CNN Settings")
             with gr.Accordion(open=False):
                 cnn_max_channel_visual = gr.Number(value=8,precision=0,label='Max CNN Channels to Preview',interactive=True)
+                cnn_couple_channel = gr.Checkbox(value=False,label='Couple Channels into RGB')
         with gr.Column():
             gr.Markdown("# Classical Settings")
             with gr.Accordion(open=False):
@@ -439,7 +440,7 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
         
         prediction_btn.click(
             fn=predict_image,
-            inputs=[image_upload,classic_show_original,cnn_max_channel_visual],
+            inputs=[image_upload,classic_show_original,cnn_max_channel_visual,cnn_couple_channel],
             outputs=[verbose,image_preview,cnn_prediction,classical_prediction,cnn_features,classical_features]
         )
         

src/model.py

@@ -55,6 +55,9 @@ class Config:
     gabor_lambda = 10
     gabor_gamma = 0.5
 
+    # Sobel
+    sobel_ksize=3
+
 class CNNFeatureExtractor(nn.Module):
     def __init__(self,config : Config):
         super().__init__()
@@ -105,51 +108,148 @@ class CNNFeatureExtractor(nn.Module):
             out = self(x)
 
         return out
-    def visualize(self, input_image, max_channels=8,show=True):
+    def visualize(
+        self,
+        input_image,
+        max_channels=8,
+        couple=False,
+        show=True,
+        **kwargs
+    ):
         self.eval()
         device = self.get_device()
 
         if isinstance(input_image, np.ndarray):
-            x = torch.from_numpy(input_image).permute(2, 0, 1).float().unsqueeze(0).to(device)  # HWC -> CHW -> B
+            x = torch.from_numpy(input_image).permute(2, 0, 1).float().unsqueeze(0).to(device)
         elif isinstance(input_image, torch.Tensor):
             x = input_image.unsqueeze(0).to(device) if input_image.ndim == 3 else input_image.to(device)
         else:
             raise TypeError("input_image must be np.ndarray or torch.Tensor")
 
-        conv_layers = [(name, module) for name, module in self.named_modules() if isinstance(module, nn.Conv2d)]
+        conv_layers = [
+            (name, module)
+            for name, module in self.named_modules()
+            if isinstance(module, nn.Conv2d)
+        ]
+
         all_layer_images = []
 
         for name, layer in conv_layers:
             activations = []
 
             def hook_fn(module, input, output):
-                activations.append(output.cpu().detach())
+                activations.append(output.detach().cpu())
 
             handle = layer.register_forward_hook(hook_fn)
             _ = self(x)
             handle.remove()
 
-            act = activations[0][0]
-            num_channels = min(act.shape[0], max_channels)
+            act = activations[0][0]  # (C, H, W)
+            C, H, W = act.shape
+
+            # --------------------------------------------------
+            # COUPLED RGB VISUALIZATION
+            # --------------------------------------------------
+            if couple:
+                max_rgb = max_channels // 3
+                num_rgb = min(C // 3, max_rgb)
+                rem = min(C - num_rgb * 3, max_channels - num_rgb * 3)
+
+                total_tiles = num_rgb + rem
+                cols = min(4, total_tiles)
+                rows = int(np.ceil(total_tiles / cols))
+
+                fig, axes = plt.subplots(
+                    rows, cols,
+                    figsize=(3 * cols, 3 * rows)
+                )
+
+                axes = np.atleast_2d(axes)
+
+                tile_idx = 0
+
+                # ---------------------------
+                # RGB COUPLED CHANNELS
+                # ---------------------------
+                for i in range(num_rgb):
+                    r = tile_idx // cols
+                    c = tile_idx % cols
+
+                    rgb = act[i*3:(i+1)*3].clone()
+
+                    for ch in range(3):
+                        v = rgb[ch]
+                        rgb[ch] = (v - v.min()) / (v.max() - v.min() + 1e-8)
+
+                    rgb = rgb.permute(1, 2, 0).numpy()
+
+                    axes[r, c].imshow(rgb)
+                    axes[r, c].axis("off")
+                    axes[r, c].set_title(f"RGB {i*3}-{i*3+2}", fontsize=9)
 
-            fig, axes = plt.subplots(1, num_channels, figsize=(3*num_channels, 3))
-            if num_channels == 1:
-                axes = [axes]
+                    tile_idx += 1
 
-            for i in range(num_channels):
-                axes[i].imshow(act[i], cmap='gray')
-                axes[i].axis('off')
+                start = num_rgb * 3
+                for j in range(rem):
+                    r = tile_idx // cols
+                    c = tile_idx % cols
+
+                    ch = act[start + j]
+                    ch = (ch - ch.min()) / (ch.max() - ch.min() + 1e-8)
+
+                    axes[r, c].imshow(ch, cmap="gray")
+                    axes[r, c].axis("off")
+                    axes[r, c].set_title(f"Ch {start + j}", fontsize=9)
+
+                    tile_idx += 1
+
+                for idx in range(tile_idx, rows * cols):
+                    r = idx // cols
+                    c = idx % cols
+                    axes[r, c].axis("off")
+
+                fig.suptitle(f"Layer: {name} (Coupled RGB + Grayscale)", fontsize=14)
+                plt.tight_layout()
+
+            # --------------------------------------------------
+            # STANDARD GRAYSCALE VISUALIZATION
+            # --------------------------------------------------
+            else:
+                num_channels = min(C, max_channels)
+                cols = min(8, num_channels)
+                rows = int(np.ceil(num_channels / cols))
+
+                fig, axes = plt.subplots(
+                    rows, cols,
+                    figsize=(3 * cols, 3 * rows)
+                )
+
+                axes = np.atleast_2d(axes)
+
+                for idx in range(num_channels):
+                    r = idx // cols
+                    c = idx % cols
+                    axes[r, c].imshow(act[idx], cmap="gray")
+                    axes[r, c].axis("off")
+
+                for idx in range(num_channels, rows * cols):
+                    r = idx // cols
+                    c = idx % cols
+                    axes[r, c].axis("off")
+
+                fig.suptitle(f"Layer: {name}", fontsize=14)
+                plt.tight_layout()
 
-            fig.suptitle(f'Layer: {name}', fontsize=14)
             if show:
                 plt.show()
 
             buf = io.BytesIO()
-            fig.savefig(buf, format='png')
+            fig.savefig(buf, format="png", dpi=150, bbox_inches="tight")
             buf.seek(0)
             img = Image.open(buf).convert("RGB")
             all_layer_images.append(np.array(img))
-            plt.close(fig) 
+            plt.close(fig)
+
         return all_layer_images
     
 class ClassicalFeatureExtractor(nn.Module):
@@ -159,74 +259,204 @@ class ClassicalFeatureExtractor(nn.Module):
         self.hog_orientations = config.hog_orientations
         self.num_downsample = config.classical_downsample
         self.config = config
-        self.feature_names = ['HoG','Canny Edge','Harris Corner','Shi-Tomasi corners','LBP','Gabor Filters']
         self.device = 'cpu'
 
     def get_device(self):
         return next(self.parameters()).device if len(list(self.parameters())) > 0 else self.device
 
 
-    def extract_features(self, img):
+    def extract_features(self, img,visualize=False,**kwargs):
         cfg = self.config
 
         # Convert to grayscale
-        min_h = cfg.hog_pixels_per_cell[0] * cfg.hog_cells_per_block[0]
-        min_w = cfg.hog_pixels_per_cell[1] * cfg.hog_cells_per_block[1]
         gray = cv2.cvtColor((img*255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
 
         for _ in range(self.num_downsample):
-            h, w = gray.shape
-            if h <= min_h or w <= min_w:
-                break
             gray = cv2.pyrDown(gray)
 
         gray = cv2.GaussianBlur(gray, cfg.gaussian_ksize, sigmaX=cfg.gaussian_sigmaX, sigmaY=cfg.gaussian_sigmaY)
+        valid_H, valid_W = gray.shape[:2]
+        
+        def render_subplots(items, max_cols=8, figsize_per_cell=3):
+            n = len(items)
+            cols = min(max_cols, n)
+            rows = int(np.ceil(n / cols))
+
+            fig, axes = plt.subplots(
+                rows, cols,
+                figsize=(cols * figsize_per_cell, rows * figsize_per_cell)
+            )
 
-        feature_list = []
+            axes = np.atleast_2d(axes)
+
+            for idx, (img, title, cmap) in enumerate(items):
+                r = idx // cols
+                c = idx % cols
+                ax = axes[r, c]
+                ax.imshow(img, cmap=cmap)
+                ax.set_title(title, fontsize=9)
+                ax.axis("off")
 
+            # Hide unused axes
+            for idx in range(n, rows * cols):
+                r = idx // cols
+                c = idx % cols
+                axes[r, c].axis("off")
+
+            plt.tight_layout()
+            return fig
+
+        feature_list = []
+        vis_items=[]
+        # figs = []
+        H, W = gray.shape
+        cell_h, cell_w = cfg.hog_pixels_per_cell
+        block_h, block_w = cfg.hog_cells_per_block
+
+        min_h = cell_h * block_h
+        min_w = cell_w * block_w
+        use_hog = (H > 2*min_h) and (W > 2*min_w)
         # 1. HOG
-        _, hog_image = hog(
-            gray,
-            orientations=cfg.hog_orientations,
-            pixels_per_cell=cfg.hog_pixels_per_cell,
-            cells_per_block=cfg.hog_cells_per_block,
-            block_norm=cfg.hog_block_norm,
-            visualize=True
-        )
-        feature_list.append(hog_image)
+        if use_hog:
+            hog_descriptors, hog_image = hog(
+                gray,
+                orientations=cfg.hog_orientations,
+                pixels_per_cell=cfg.hog_pixels_per_cell,
+                cells_per_block=cfg.hog_cells_per_block,
+                block_norm=cfg.hog_block_norm,
+                visualize=True,
+                feature_vector=False
+            )
 
+            hog_cells = hog_descriptors.mean(axis=(2, 3))
+            
+            cell_h, cell_w = cfg.hog_pixels_per_cell
+            hog_pixel = np.repeat(
+                np.repeat(hog_cells, cell_h, axis=0),
+                cell_w, axis=1
+            )
+            hog_pixel = hog_pixel[:gray.shape[0], :gray.shape[1]]
+            hog_energy = np.sum(hog_pixel, axis=2)
+            dominant_bin = np.argmax(hog_pixel, axis=2)
+            dominant_strength = np.max(hog_pixel, axis=2)
+            dominant_weighted = dominant_bin * dominant_strength
+            valid_H, valid_W = hog_pixel.shape[:2]
+            if visualize:
+                # figs.append(plot_feature(hog_energy, "HOG Energy"))
+                # figs.append(plot_feature(dominant_bin, "HOG Dominant Bin",cmap='hsv'))
+                # figs.append(plot_feature(dominant_weighted, "HOG Weighted Dominant Bin"))
+                # figs.append(plot_feature(hog_image[:valid_H, :valid_W], f"HoG"))
+                vis_items.append((hog_energy, "HOG Energy",'gray'))
+                vis_items.append((dominant_bin, "HOG Dominant Bin",'hsv'))
+                vis_items.append((dominant_weighted, "HOG Weighted Dominant Bin",'gray'))
+                vis_items.append((hog_image[:valid_H, :valid_W], f"HoG",'gray'))
+            for b in range(hog_pixel.shape[2]):
+                feature_list.append(hog_pixel[:, :, b])
+        
+        
         # 2. Canny edges
         edges = cv2.Canny(gray, cfg.canny_low, cfg.canny_high) / 255.0
-        feature_list.append(edges)
-
+        # feature_list.append(edges.ravel())
+        feature_list.append(edges[:valid_H, :valid_W])
+        if visualize:
+            # figs.append(plot_feature(edges[:valid_H, :valid_W], "Canny Edge"))
+            vis_items.append((edges[:valid_H, :valid_W], "Canny Edge", "gray"))
         # 3. Harris corners
         harris = cv2.cornerHarris(gray, blockSize=cfg.harris_block_size, ksize=cfg.harris_ksize, k=cfg.harris_k)
         harris = cv2.dilate(harris, None)
         harris = np.clip(harris, 0, 1)
-        feature_list.append(harris)
-
-        # 4. Shi-Tomasi corners
-        shi_corners = np.zeros_like(gray, dtype=np.float32)
-        keypoints = cv2.goodFeaturesToTrack(gray, maxCorners=cfg.shi_max_corners, qualityLevel=cfg.shi_quality_level, minDistance=cfg.shi_min_distance)
-        if keypoints is not None:
-            for kp in keypoints:
-                x, y = kp.ravel()
-                shi_corners[int(y), int(x)] = 1.0
-        feature_list.append(shi_corners)
-
+        # feature_list.append(harris.ravel())
+        feature_list.append(harris[:valid_H, :valid_W])
+        if visualize:
+            # figs.append(plot_feature(harris[:valid_H, :valid_W], "Harris Corner"))
+            vis_items.append((harris[:valid_H, :valid_W], "Harris Corner", "gray"))
+        # # 4. Shi-Tomasi corners
+        # shi_corners = np.zeros_like(gray, dtype=np.float32)
+        # keypoints = cv2.goodFeaturesToTrack(gray, maxCorners=cfg.shi_max_corners, qualityLevel=cfg.shi_quality_level, minDistance=cfg.shi_min_distance)
+        # if keypoints is not None:
+        #     for kp in keypoints:
+        #         x, y = kp.ravel()
+        #         shi_corners[int(y), int(x)] = 1.0
+        # # feature_list.append(shi_corners.ravel())
+        # feature_list.append(shi_corners[:valid_H, :valid_W])
+        # if visualize:
+        #     figs.append(plot_feature(shi_corners[:valid_H, :valid_W], "Shi-Tomasi Corner"))
         # 5. LBP
         lbp = local_binary_pattern(gray, P=cfg.lbp_P, R=cfg.lbp_R, method='uniform')
         lbp = lbp / lbp.max() if lbp.max() != 0 else lbp
-        feature_list.append(lbp)
-
+        # feature_list.append(lbp.ravel())
+        feature_list.append(lbp[:valid_H, :valid_W])
+        if visualize:
+            # figs.append(plot_feature(lbp[:valid_H, :valid_W], "LBP"))
+            vis_items.append((lbp[:valid_H, :valid_W], "LBP", "gray"))
         # 6. Gabor filter
-        g_kernel = cv2.getGaborKernel((cfg.gabor_ksize, cfg.gabor_ksize), cfg.gabor_sigma, cfg.gabor_theta, cfg.gabor_lambda, cfg.gabor_gamma)
-        gabor_feat = cv2.filter2D(gray, cv2.CV_32F, g_kernel)
-        gabor_feat = (gabor_feat - gabor_feat.min()) / (gabor_feat.max() - gabor_feat.min() + 1e-8)
-        feature_list.append(gabor_feat)
+        # g_kernel = cv2.getGaborKernel((cfg.gabor_ksize, cfg.gabor_ksize), cfg.gabor_sigma, cfg.gabor_theta, cfg.gabor_lambda, cfg.gabor_gamma)
+        # gabor_feat = cv2.filter2D(gray, cv2.CV_32F, g_kernel)
+        # gabor_feat = (gabor_feat - gabor_feat.min()) / (gabor_feat.max() - gabor_feat.min() + 1e-8)
+        # # feature_list.append(gabor_feat.ravel())
+        # feature_list.append(gabor_feat[:valid_H, :valid_W])
+        # if visualize:
+        #     figs.append(plot_feature(gabor_feat[:valid_H, :valid_W], "Gabor Filter"))
+
+        for theta in [0, np.pi/4, np.pi/2]:
+            kernel = cv2.getGaborKernel(
+                (cfg.gabor_ksize, cfg.gabor_ksize),
+                cfg.gabor_sigma, theta,
+                cfg.gabor_lambda, cfg.gabor_gamma
+            )
+            g = cv2.filter2D(gray, cv2.CV_32F, kernel)
+            g = np.abs(g)
+            g /= g.max() + 1e-8
+            feature_list.append(g[:valid_H, :valid_W])
+            if visualize:
+                # figs.append(plot_feature(g[:valid_H, :valid_W], "Gabor Filter"))
+                vis_items.append((g[:valid_H, :valid_W], f"Gabor θ={theta:.2f}", "gray"))
+        # 7. Sobel
+        sobelx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=cfg.sobel_ksize)
+        sobely = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=cfg.sobel_ksize)
+
+        sobelx = np.abs(sobelx)
+        sobely = np.abs(sobely)
+
+        sobelx /= sobelx.max() + 1e-8
+        sobely /= sobely.max() + 1e-8
+
+        feature_list.append(sobelx[:valid_H, :valid_W])
+        feature_list.append(sobely[:valid_H, :valid_W])
+        if visualize:
+            # figs.append(plot_feature(sobelx[:valid_H, :valid_W], "Sobel X"))
+            # figs.append(plot_feature(sobely[:valid_H, :valid_W], "Sobel Y"))
+            vis_items.append((sobelx[:valid_H, :valid_W], "Sobel X",'gray'))
+            vis_items.append((sobely[:valid_H, :valid_W], "Sobel Y",'gray'))
+        # 8. Laplacian
+        lap = cv2.Laplacian(gray, cv2.CV_32F)
+        lap = np.abs(lap)
+        lap /= lap.max() + 1e-8
+
+        feature_list.append(lap[:valid_H, :valid_W])
+
+        if visualize:
+            # figs.append(plot_feature(lap[:valid_H, :valid_W], "Laplacian"))
+            vis_items.append((lap[:valid_H, :valid_W], "Laplacian",'gray'))
+
+        # 9. Gradient Magnitude
+        gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=cfg.sobel_ksize)
+        gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=cfg.sobel_ksize)
+
+        grad_mag = np.sqrt(gx**2 + gy**2)
+        grad_mag /= grad_mag.max() + 1e-8
+
+        feature_list.append(grad_mag[:valid_H, :valid_W])
+
+        if visualize:
+            # figs.append(plot_feature(grad_mag[:valid_H, :valid_W], "Gradient Magnitude"))
+            vis_items.append((grad_mag[:valid_H, :valid_W], "Gradient Magnitude",'gray'))
 
         # Stack all features along channel axis
-        features = np.stack(feature_list, axis=2)
+        features = np.stack(feature_list, axis=0)
+        # features = np.concatenate(feature_list).astype(np.float32)
+        if visualize:
+            return features.astype(np.float32),[render_subplots(vis_items, max_cols=8)]
         return features.astype(np.float32)
 
 
@@ -248,15 +478,13 @@ class ClassicalFeatureExtractor(nn.Module):
             feat = self.extract_features(img)
             batch_features.append(feat)
         batch_features = np.stack(batch_features, axis=0)
-        return torch.from_numpy(batch_features).float().to(self.get_device())
+        batch_features = torch.from_numpy(batch_features).float().to(self.get_device())
+        return batch_features
     
     def visualize(self, img, show_original=True,show=True):
         if img.ndim != 3 or img.shape[2] != 3:
             img = np.repeat(img[:, :, None], 3, axis=2)
 
-        feature_stack = self.extract_features(img)
-        num_channels = feature_stack.shape[2]
-
         outputs = []  
 
         def fig_to_pil(fig):
@@ -279,15 +507,10 @@ class ClassicalFeatureExtractor(nn.Module):
             if show:
                 plt.show()                      
             outputs.append(fig_to_pil(fig)) 
-
-        for c in range(num_channels):
-            fig = plt.figure(figsize=(4, 4))
-
-            plt.imshow(feature_stack[:, :, c], cmap="gray")
-            plt.title(f"Feature {self.feature_names[c]}")
-            plt.axis("off")
-            if show:
-                plt.show()                      
+        feature_stack,figs = self.extract_features(img,visualize=True)
+        if show:
+            plt.show()      
+        for fig in figs:
             outputs.append(fig_to_pil(fig)) 
 
         return outputs