Update

app.py

@@ -2,10 +2,10 @@ import gradio as gr
 import zipfile
 import os
 import torch
-from src.dataloader import ImageDataset,collate_fn
+from src.dataloader import ImageDataset,collate_fn,AugmentedSubset,simple_augment
 from src.model import Classifier,Config,CNNFeatureExtractor,ClassicalFeatureExtractor,load
 from torch.utils.data import Subset
-from src.trainer import ModelTrainer
+from src.trainer import ModelTrainer,model_evaluation
 import torch
 import os
 import numpy as np
@@ -17,9 +17,6 @@ import matplotlib.pyplot as plt
 import shutil
 import pandas as pd
 from sklearn.model_selection import train_test_split
-from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
-from sklearn.metrics import classification_report
-from torch.utils.data import DataLoader
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 def unzip_dataset(zip_file):
@@ -82,13 +79,22 @@ def plot(datas, labels, xlabel, ylabel, title, figsize=(16, 8)):
 
 class TrainingInterrupted(Exception):
     pass
+cnntrainer=None
+classictrainer=None
 def stop_training():
     global training_interrupt
     training_interrupt = True
+    if cnntrainer is not None:
+        cnntrainer.interrupt=True
+    if classictrainer is not None:
+        classictrainer.interrupt=True
     return "Training stopped."
 
+
+
 def train(cnn,classic,train_set,val_set,batch_size,lr,epochs,device="cpu",visualize_every=5):
     global training_interrupt
+    global cnntrainer,classictrainer
     training_interrupt = False
     global cnn_history
     global classic_history
@@ -163,17 +169,17 @@ def train(cnn,classic,train_set,val_set,batch_size,lr,epochs,device="cpu",visual
         print(e)
         return
 
-def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,
+def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,use_augment,
                 img_width,img_height,fc_num_layers,
                 in_channels,conv_hidden_dim,dropout,
                 classical_downsample,
-                hog_orientations,hog_pixels_per_cell,hog_cells_per_block,hog_block_norm,
+                # hog_orientations,hog_pixels_per_cell,hog_cells_per_block,hog_block_norm,
                 canny_sigma,canny_low,canny_high,
                 gaussian_ksize,gaussian_sigmaX,gaussian_sigmaY,
                 harris_block_size,harris_ksize,harris_k,
-                shi_max_corners,shi_quality_level,shi_min_distance,
                 lbp_P,lbp_R,
-                gabor_ksize,gabor_sigma,gabor_theta,gabor_lambda,gabor_gamma):
+                gabor_ksize,gabor_sigma,gabor_theta,gabor_lambda,gabor_gamma,
+                sobel_ksize):
     config = Config()
     global training_interrupt
     training_interrupt = False
@@ -190,10 +196,10 @@ def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,
     config.dropout=dropout
     # Classical Config
     config.classical_downsample=int(classical_downsample)
-    config.hog_orientations=int(hog_orientations)
-    config.hog_pixels_per_cell=(int(hog_pixels_per_cell),int(hog_pixels_per_cell))
-    config.hog_cells_per_block=(int(hog_cells_per_block),int(hog_cells_per_block))
-    config.hog_block_norm=hog_block_norm
+    # config.hog_orientations=int(hog_orientations)
+    # config.hog_pixels_per_cell=(int(hog_pixels_per_cell),int(hog_pixels_per_cell))
+    # config.hog_cells_per_block=(int(hog_cells_per_block),int(hog_cells_per_block))
+    # config.hog_block_norm=hog_block_norm
     config.canny_sigma=int(canny_sigma)
     config.canny_low=canny_low
     config.canny_high=canny_high
@@ -203,9 +209,6 @@ def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,
     config.harris_block_size=int(harris_block_size)
     config.harris_ksize=int(harris_ksize)
     config.harris_k=harris_k
-    config.shi_max_corners=int(shi_max_corners)
-    config.shi_quality_level=shi_quality_level
-    config.shi_min_distance=int(shi_min_distance)
     config.lbp_P=int(lbp_P)
     config.lbp_R=int(lbp_R)
     config.gabor_ksize=int(gabor_ksize)
@@ -213,12 +216,14 @@ def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,
     config.gabor_theta=int(gabor_theta)
     config.gabor_lambda=int(gabor_lambda)
     config.gabor_gamma=gabor_gamma
+    config.sobel_ksize=int(sobel_ksize)
     cnn_history_plots=[]
     classical_history_plots=[]
     cnn_plotted=False
     try:
         dataset = ImageDataset(DATASET_PATH,config.img_size)
         labels = [item['id'] for item in dataset.data]
+        classes_name = dataset.classes
         train_idx, validation_idx = train_test_split(np.arange(len(dataset)),
                                                 test_size=0.2,
                                                 random_state=SEED,
@@ -226,6 +231,12 @@ def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,
                                                 stratify=labels)
         train_dataset = Subset(dataset, train_idx)
         val_dataset = Subset(dataset, validation_idx)
+        if use_augment:
+            train_dataset = AugmentedSubset(train_dataset,simple_augment)
+            val_dataset = AugmentedSubset(val_dataset,None)
+        else:
+            train_dataset = AugmentedSubset(train_dataset,None)
+            val_dataset = AugmentedSubset(val_dataset,None)
         cnnbackbone = CNNFeatureExtractor(config).to(device)
         cnnmodel = Classifier(cnnbackbone,train_dataset.dataset.classes,config).to(device)
         classicbackbone = ClassicalFeatureExtractor(config)
@@ -235,11 +246,17 @@ def train_model(zip_file,batch_size,lr,epochs,seed,vis_every,
                 cnn_plotted=True
                 cnn_history_plots.append(plot([cnn_history['train_acc'],cnn_history['val_acc']],['Training Accuracy','Validation Accuracy'],'Epochs','Accuracy (%)','Training vs Validation Accuracy'))
                 cnn_history_plots.append(plot([cnn_history['train_loss'],cnn_history['val_loss']],['Training Loss','Validation Loss'],'Epochs','Loss','Training vs Validation Loss'))
-
+                cm,cr,roc = model_evaluation(cnnmodel,val_dataset,device,BATCH_SIZE,0,classes_name)
+                cnn_history_plots.append(fig_to_image(cm))
+                cnn_history_plots.append(fig_to_image(cr))
+                cnn_history_plots.append(fig_to_image(roc))
             yield cnn_text,cnn_fig,classic_text,classic_fig,cnn_history_plots,classical_history_plots
         classical_history_plots.append(plot([classic_history['train_acc'],classic_history['val_acc']],['Training Accuracy','Validation Accuracy'],'Epochs','Accuracy (%)','Training vs Validation Accuracy'))
         classical_history_plots.append(plot([classic_history['train_loss'],classic_history['val_loss']],['Training Loss','Validation Loss'],'Epochs','Loss','Training vs Validation Loss'))
-
+        cm,cr,roc = model_evaluation(classicmodel,val_dataset,device,BATCH_SIZE,0,classes_name)
+        classical_history_plots.append(fig_to_image(cm))
+        classical_history_plots.append(fig_to_image(cr))
+        classical_history_plots.append(fig_to_image(roc))
         yield cnn_text,cnn_fig,classic_text,classic_fig,cnn_history_plots,classical_history_plots
 
     except RuntimeError as e:
@@ -314,6 +331,7 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
             epochs= gr.Number(value=20,label="Epochs",interactive=True,precision=0)
             seed=gr.Number(value=42,label='Seed',interactive=True,precision=0)
             vis_every=gr.Number(value=5,label='Visualize Validation Every (Epochs)',interactive=True,precision=0)
+            use_augment = gr.Checkbox(value=True,label='Use data augmentation for train data')
         with gr.Row():
             img_width=gr.Number(value=128,label='Image Width',interactive=True,precision=0)
             img_height=gr.Number(value=128,label='Image Height',interactive=True,precision=0)
@@ -328,11 +346,12 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
         with gr.Accordion(label='Classical Feature Extractor Settings',open=False):
             with gr.Row():
                 classical_downsample = gr.Number(value=1,label='Classical Extractor Downsampling Amount',interactive=True,precision=0)
-            with gr.Row():
-                hog_orientations = gr.Number(value=9,label='HoG Orientations',interactive=True,precision=0)
-                hog_pixels_per_cell = gr.Number(value=16,label='HoG pixels per cell',interactive=True,precision=0)
-                hog_cells_per_block = gr.Number(value=2,label='HoG cells per block',interactive=True,precision=0)
-                hog_block_norm = gr.Dropdown(['L2-Hys'],value='L2-Hys',label='HoG Block Normalization Method',interactive=True)
+            # Deprecated
+            # with gr.Row():
+            #     hog_orientations = gr.Number(value=9,label='HoG Orientations',interactive=True,precision=0)
+            #     hog_pixels_per_cell = gr.Number(value=16,label='HoG pixels per cell',interactive=True,precision=0)
+            #     hog_cells_per_block = gr.Number(value=2,label='HoG cells per block',interactive=True,precision=0)
+            #     hog_block_norm = gr.Dropdown(['L2-Hys'],value='L2-Hys',label='HoG Block Normalization Method',interactive=True)
             with gr.Row():
                 canny_sigma = gr.Number(value=1.0,label='Canny Sigma Value',interactive=True)
                 canny_low = gr.Number(value=100,label='Canny Low Threshold',interactive=True,precision=0)
@@ -345,19 +364,17 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
                 harris_block_size = gr.Number(value=2,label='Harris Corner Block Size',interactive=True,precision=0)
                 harris_ksize = gr.Number(value=3,label='Harris Corner Kernel Size',interactive=True,precision=0)
                 harris_k = gr.Slider(minimum=0.01, maximum=0.1, value=0.04, step=0.005, label='Harris Corner K value',interactive=True)
-            with gr.Row():
-                shi_max_corners = gr.Number(value=100,label='Shi-Tomasi Max Corners',interactive=True,precision=0)
-                shi_quality_level = gr.Number(value=0.01,label='Shi-Tomasi Quality Level',interactive=True)
-                shi_min_distance = gr.Number(value=10,label='Shi-Tomasi Min Distance',interactive=True,precision=0)
             with gr.Row():
                 lbp_P = gr.Number(value=8,label='LBP P Value',interactive=True,precision=0)
                 lbp_R = gr.Number(value=1,label='LBP R Value',interactive=True,precision=0)
             with gr.Row():
                 gabor_ksize  = gr.Number(value=21,label="Gabor Kernel Size",interactive=True,precision=0)
                 gabor_sigma  = gr.Number(value=5,label="Gabor Sigma",interactive=True,precision=0)
-                gabor_theta  = gr.Number(value=0,label="Gabor Theta",interactive=True,precision=0)
+                gabor_theta  = gr.Number(value=0,label="Gabor Theta",interactive=True,precision=0,info="This current does nothing")
                 gabor_lambda = gr.Number(value=10,label="Gabor Lambda",interactive=True,precision=0)
                 gabor_gamma  = gr.Number(value=0.5,label="Gabor Gamma",interactive=True)
+            with gr.Row():
+                sobel_ksize = gr.Number(value=3,label="Sobel Kernel Size",interactive=True,precision=0)
         with gr.Column():
             train_btn = gr.Button("Train Model",variant='secondary',interactive=True)
             stop_btn = gr.Button("Stop Training")
@@ -375,17 +392,17 @@ with gr.Blocks(title="Object Classifier Playground") as demo:
                 classical_plots = gr.Gallery(label="CNN Training Performance",interactive=False,object_fit='fill',columns=1)
         stop_btn.click(fn=stop_training, inputs=[], outputs=[])
         train_btn.click(fn=train_model,
-                        inputs=[zip_file,batch_size,lr,epochs,seed,vis_every,
+                        inputs=[zip_file,batch_size,lr,epochs,seed,vis_every,use_augment,
                             img_width,img_height,fc_num_layers,
                             in_channels,conv_hidden_dim,dropout,
                             classical_downsample,
-                            hog_orientations,hog_pixels_per_cell,hog_cells_per_block,hog_block_norm,
+                            # hog_orientations,hog_pixels_per_cell,hog_cells_per_block,hog_block_norm,
                             canny_sigma,canny_low,canny_high,
                             gaussian_ksize,gaussian_sigmaX,gaussian_sigmaY,
                             harris_block_size,harris_ksize,harris_k,
-                            shi_max_corners,shi_quality_level,shi_min_distance,
                             lbp_P,lbp_R,
-                            gabor_ksize,gabor_sigma,gabor_theta,gabor_lambda,gabor_gamma],
+                            gabor_ksize,gabor_sigma,gabor_theta,gabor_lambda,gabor_gamma,
+                            sobel_ksize],
                         outputs=[cnn_log,cnn_fig,classical_log,classical_fig,cnn_plots,classical_plots]
                         )
     def make_figure_from_image(img):

src/dataloader.py

@@ -1,4 +1,4 @@
-from torch.utils.data import Dataset
+from torch.utils.data import Subset,Dataset
 import torch
 import os
 import numpy as np
@@ -37,4 +37,25 @@ class ImageDataset(Dataset):
         img = img.astype(np.float32) / 255.0
         return img,label
 
+def simple_augment(img):
+    if np.random.rand() > 0.5:
+        img = cv2.flip(img, 1)
 
+    angle = np.random.uniform(-15, 15)
+    h, w = img.shape[:2]
+    M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)
+    img = cv2.warpAffine(img, M, (w, h), borderMode=cv2.BORDER_REFLECT)
+
+    return img
+
+
+class AugmentedSubset(Subset):
+    def __init__(self, subset, augment_fn=None):
+        super().__init__(subset.dataset, subset.indices)
+        self.augment_fn = augment_fn
+
+    def __getitem__(self, idx):
+        img, label = super().__getitem__(idx)
+        if self.augment_fn:
+            img = self.augment_fn(img)
+        return img, label

src/model.py

@@ -4,6 +4,8 @@ import cv2
 import numpy as np
 from dataclasses import dataclass
 from skimage.feature import hog,local_binary_pattern
+import itertools
+import torch.nn.functional as F
 import matplotlib.pyplot as plt
 import os
 import io
@@ -14,7 +16,7 @@ class Config:
     img_size=(256,256)
     in_channels=3
     fc_num_layers=3
-    conv_hidden_dim=3
+    conv_hidden_dim=2
     conv_kernel_size=3
     dropout=0.2
     classical_downsample=1
@@ -39,10 +41,6 @@ class Config:
     harris_ksize = 3
     harris_k = 0.04
 
-    # Shi-Tomasi corners
-    shi_max_corners = 100
-    shi_quality_level = 0.01
-    shi_min_distance = 10
 
     # LBP
     lbp_P = 8 
@@ -58,6 +56,7 @@ class Config:
     # Sobel
     sobel_ksize=3
 
+
 class CNNFeatureExtractor(nn.Module):
     def __init__(self,config : Config):
         super().__init__()
@@ -67,16 +66,16 @@ class CNNFeatureExtractor(nn.Module):
         self.img_size = config.img_size
         out_channel = 32
         for i in range(config.conv_hidden_dim):
-            layers.append(nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=config.conv_kernel_size,stride=1,padding=1))
+            layers.append(nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=config.conv_kernel_size,stride=1,padding=config.conv_kernel_size // 2))
             layers.append(nn.BatchNorm2d(out_channel))
             layers.append(nn.ReLU())
-            layers.append(nn.MaxPool2d(2))
+            layers.append(nn.MaxPool2d((2,2)))
             in_channel=out_channel
             out_channel*=2
         self.layers = nn.Sequential(*layers)
     def get_device(self):
         return next(self.parameters()).device
-    def forward(self,x):
+    def forward(self,x,**kwargs):
         if isinstance(x, list):
             if isinstance(x[0], np.ndarray):
                 x = np.stack(x, axis=0) 
@@ -129,7 +128,7 @@ class CNNFeatureExtractor(nn.Module):
         conv_layers = [
             (name, module)
             for name, module in self.named_modules()
-            if isinstance(module, nn.Conv2d)
+            if isinstance(module, nn.ReLU)
         ]
 
         all_layer_images = []
@@ -251,7 +250,7 @@ class CNNFeatureExtractor(nn.Module):
             plt.close(fig)
 
         return all_layer_images
-    
+
 class ClassicalFeatureExtractor(nn.Module):
     def __init__(self, config : Config):
         super().__init__()
@@ -260,128 +259,103 @@ class ClassicalFeatureExtractor(nn.Module):
         self.num_downsample = config.classical_downsample
         self.config = config
         self.device = 'cpu'
-
+        self.convolution=None
+    
     def get_device(self):
         return next(self.parameters()).device if len(list(self.parameters())) > 0 else self.device
 
-
+    def render_subplots(self,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)
+        )
+        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")
+        for idx in range(n, rows * cols):
+            r = idx // cols
+            c = idx % cols
+            axes[r, c].axis("off")
+
+        plt.tight_layout()
+        return fig
+    
     def extract_features(self, img,visualize=False,**kwargs):
         cfg = self.config
-
         # Convert to grayscale
         gray = cv2.cvtColor((img*255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
-
         for _ in range(self.num_downsample):
             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)
-            )
-
-            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
-        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))
+        # DEPRECATED
+        # 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 = False
+        # # 1. HOG
+        # 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])
+        #     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:
+        #         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.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.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
+
+        # 4. 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.ravel())
@@ -389,15 +363,7 @@ class ClassicalFeatureExtractor(nn.Module):
         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.ravel())
-        # feature_list.append(gabor_feat[:valid_H, :valid_W])
-        # if visualize:
-        #     figs.append(plot_feature(gabor_feat[:valid_H, :valid_W], "Gabor Filter"))
-
+        # 5. Gabor filter
         for theta in [0, np.pi/4, np.pi/2]:
             kernel = cv2.getGaborKernel(
                 (cfg.gabor_ksize, cfg.gabor_ksize),
@@ -409,9 +375,8 @@ class ClassicalFeatureExtractor(nn.Module):
             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
+        # 6. 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)
 
@@ -424,11 +389,9 @@ class ClassicalFeatureExtractor(nn.Module):
         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
+        # 7. Laplacian
         lap = cv2.Laplacian(gray, cv2.CV_32F)
         lap = np.abs(lap)
         lap /= lap.max() + 1e-8
@@ -436,10 +399,9 @@ class ClassicalFeatureExtractor(nn.Module):
         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
+        # 8. 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)
 
@@ -449,37 +411,38 @@ class ClassicalFeatureExtractor(nn.Module):
         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=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),[self.render_subplots(vis_items, max_cols=8)]
         return features.astype(np.float32)
 
 
-    def forward(self, x):
+    def forward(self, x, **kwargs):
+        if isinstance(x, list):
+            x = np.stack(x, axis=0)
+
         if isinstance(x, torch.Tensor):
             x = x.cpu().numpy()
+
         if isinstance(x, np.ndarray):
-            if x.ndim == 3:
-                x = np.expand_dims(x, 0)
+            if x.ndim == 3:         
+                x = x[None]
             elif x.ndim != 4:
-                raise ValueError(f"Expected input of shape HWC or BHWC, got {x.shape}")
-        elif isinstance(x, list):
-            x = np.stack(x, axis=0)
-
-        batch_features = []
+                raise ValueError(
+                    f"Expected input of shape HWC or BHWC, got {x.shape}"
+                )
+        feats = []
         for img in x:
-            if img.ndim != 3 or img.shape[2] != 3:
+            if img.shape[2] != 3:
                 img = np.repeat(img[:, :, None], 3, axis=2)
-            feat = self.extract_features(img)
-            batch_features.append(feat)
-        batch_features = np.stack(batch_features, axis=0)
-        batch_features = torch.from_numpy(batch_features).float().to(self.get_device())
-        return batch_features
+            feats.append(self.extract_features(img))
+
+        feats = np.stack(feats, axis=0)
+        feats = torch.from_numpy(feats).float().to(self.get_device())
+        return feats
     
     def visualize(self, img, show_original=True,show=True):
         if img.ndim != 3 or img.shape[2] != 3:
@@ -517,10 +480,14 @@ class ClassicalFeatureExtractor(nn.Module):
 
 
     def output(self):
-        """Return dummy output to compute in_features for FC head"""
-        dummy_img = np.zeros((1, self.img_size[1],self.img_size[0], 3), dtype=np.float32)
-        feat = self.forward(dummy_img)
-        return feat
+        dummy = np.zeros(
+            (self.img_size[1], self.img_size[0], 3),
+            dtype=np.float32
+        )
+
+        feats = self.forward(dummy)
+
+        return feats
 
 
 
@@ -530,21 +497,20 @@ class FullyConnectedHead(nn.Module):
         num_classes = len(classes)
         self.classes = classes
         layers = []
-        out_features=256
-        for i in range(config.fc_num_layers):
-            layers.append(nn.Linear(in_features,out_features))
-            layers.append(nn.BatchNorm1d(out_features))
+        hidden_dim =1024
+        for _ in range(config.fc_num_layers):
+            layers.append(nn.Linear(in_features, hidden_dim))
+            layers.append(nn.BatchNorm1d(hidden_dim))
             layers.append(nn.ReLU())
             layers.append(nn.Dropout(config.dropout))
-            in_features=out_features
-            out_features=out_features // 2
-            if out_features <= num_classes:
-                break
+
+            in_features = hidden_dim
+            hidden_dim = max(hidden_dim // 2, num_classes * 2)
         layers.append(nn.Linear(in_features,num_classes))
         self.layers = nn.Sequential(*layers)
     def get_device(self):
         return next(self.parameters()).device
-    def forward(self,x : torch.Tensor):
+    def forward(self,x : torch.Tensor,**kwargs):
         x=x.to(self.get_device())
         return self.layers(x)
     
@@ -568,15 +534,15 @@ class Classifier(nn.Module):
         target_size = self.config.img_size
         x = cv2.resize(x, target_size)
         logits = self.forward(x)    
-        probs = torch.softmax(logits, dim=1)
+        probs = torch.softmax(logits,dim=1)
         pred_idx = torch.argmax(probs, dim=1).item()
 
         return self.classes[pred_idx]
 
-    def forward(self,x):
-        feat = self.backbone(x)
-        feat = self.flatten(feat)
-        return self.head(feat)
+    def forward(self,x,**kwargs):
+        feat = self.backbone(x,**kwargs)
+        feat = self.flatten(feat,**kwargs)
+        return self.head(feat,**kwargs)
     def visualize_feature(self,img,return_img=True,**kwargs):
         target_size = self.config.img_size
         img = cv2.resize(img, target_size)

src/trainer.py

@@ -10,12 +10,149 @@ import random
 import numpy as np
 import torch.nn as nn
 import time
-
+from sklearn.metrics import (
+    confusion_matrix,
+    classification_report,
+    roc_curve,
+    auc
+)
+from sklearn.preprocessing import label_binarize
 def seed_worker(worker_id):
     worker_seed = torch.initial_seed() % 2**32
     np.random.seed(worker_seed)
     random.seed(worker_seed)
 
+def model_evaluation(model, val_set, device,batch_size=32,num_workers=0, class_names=None):
+
+    model.eval()
+    all_preds = []
+    all_probs = []
+    all_labels = []
+    val_loader = DataLoader(
+        val_set,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers
+    )
+    with torch.no_grad():
+        for images, labels in val_loader:
+            if images.ndim == 4 and images.shape[-1] in (1, 3):
+                images = images.permute(0, 3, 1, 2)
+            images = images.to(device)
+            labels = labels.to(device)
+            logits = model(images)
+            probs = torch.softmax(logits, dim=1)
+            preds = torch.argmax(probs, dim=1)
+
+            all_preds.append(preds.cpu().numpy())
+            all_probs.append(probs.cpu().numpy())
+            all_labels.append(labels.cpu().numpy())
+
+    y_true = np.concatenate(all_labels)
+    y_pred = np.concatenate(all_preds)
+    y_prob = np.concatenate(all_probs)
+
+    num_classes = y_prob.shape[1]
+
+    if class_names is None:
+        class_names = [f"Class {i}" for i in range(num_classes)]
+
+    cm = confusion_matrix(y_true, y_pred)
+
+    cm_fig, ax = plt.subplots(figsize=(6, 6))
+    im = ax.imshow(cm)
+
+    ax.set_title("Confusion Matrix")
+    ax.set_xlabel("Predicted")
+    ax.set_ylabel("True")
+    ax.set_xticks(range(num_classes))
+    ax.set_yticks(range(num_classes))
+    ax.set_xticklabels(class_names, rotation=75)
+    ax.set_yticklabels(class_names)
+
+    for i in range(num_classes):
+        for j in range(num_classes):
+            ax.text(j, i, cm[i, j], ha="center", va="center")
+
+    plt.tight_layout()
+
+    report = classification_report(
+        y_true, y_pred,
+        target_names=class_names,
+        output_dict=True
+    )
+
+    cr_fig, ax = plt.subplots(figsize=(12, 8))
+    ax.axis("off")
+
+    table_data = []
+    headers = ["Class", "Precision", "Recall", "F1", "Support"]
+
+    for cls in class_names:
+        row = report[cls]
+        table_data.append([
+            cls,
+            f"{row['precision']:.3f}",
+            f"{row['recall']:.3f}",
+            f"{row['f1-score']:.3f}",
+            int(row['support'])
+        ])
+
+    accuracy = report["accuracy"]
+    macro_avg = report["macro avg"]
+    weighted_avg = report["weighted avg"]
+
+    table_data.append([
+        "Accuracy",
+        f"{accuracy:.3f}",
+        "",
+        "",
+        ""
+    ])
+
+    table_data.append([
+        "Macro Avg",
+        f"{macro_avg['precision']:.3f}",
+        f"{macro_avg['recall']:.3f}",
+        f"{macro_avg['f1-score']:.3f}",
+        f"{int(macro_avg['support'])}" if 'support' in macro_avg else ""
+    ])
+
+    table_data.append([
+        "Weighted Avg",
+        f"{weighted_avg['precision']:.3f}",
+        f"{weighted_avg['recall']:.3f}",
+        f"{weighted_avg['f1-score']:.3f}",
+        f"{int(weighted_avg['support'])}" if 'support' in weighted_avg else ""
+    ])
+
+    table = ax.table(
+        cellText=table_data,
+        colLabels=headers,
+        loc="center"
+    )
+
+    table.scale(1, 2)
+    ax.set_title("Classification Report")
+
+    y_true_bin = label_binarize(y_true, classes=list(range(num_classes)))
+
+    roc_fig, ax = plt.subplots(figsize=(6, 6))
+
+    for i in range(num_classes):
+        fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_prob[:, i])
+        roc_auc = auc(fpr, tpr)
+        ax.plot(fpr, tpr, label=f"{class_names[i]} (AUC={roc_auc:.3f})")
+
+    ax.plot([0, 1], [0, 1], linestyle="--")
+    ax.set_xlabel("False Positive Rate")
+    ax.set_ylabel("True Positive Rate")
+    ax.set_title("ROC-AUC Curve")
+    ax.legend()
+    ax.grid(True)
+
+    return cm_fig, cr_fig, roc_fig
+
 class ModelTrainer:
     def __init__(self,model : Classifier,train_set : ImageDataset,val_set : ImageDataset = None, batch_size=32,lr = 1e-3,device='cpu',return_fig=False, seed=None):
         g = torch.Generator()
@@ -50,6 +187,9 @@ class ModelTrainer:
         self.optim.zero_grad()
         self.criterion = nn.CrossEntropyLoss()
         self.return_fig=return_fig
+        self.best_model_state = None
+        self.best_val_acc = 0.0
+        self.interrupt=False
     
     def visualize_batch(self, imgs, preds, labels, class_names=None, max_samples=4):
         
@@ -114,13 +254,15 @@ class ModelTrainer:
             return None
 
 
-    def train_one_epoch(self):
+    def train_one_epoch(self,epoch):
         self.model.train()
         total_loss = 0
         train_pbar = tqdm(self.train_loader, desc="Training",leave=False)
         correct = 0
         total = 0
         for imgs, labels in train_pbar:
+            if self.interrupt:
+                break
             labels = labels.to(self.device)
 
             # Forward
@@ -146,18 +288,30 @@ class ModelTrainer:
         val_losses=[]
         val_accuracies=[]
         for epoch in range(1, epochs + 1):
-            train_loss,train_acc = self.train_one_epoch()
+            train_loss,train_acc = self.train_one_epoch(epoch)
+            if self.interrupt:
+                return
             train_losses.append(train_loss)
             train_accuracies.append(train_acc)
             if self.val_loader is not None:
                 val_loss,val_acc,fig=self.validate(epoch, visualize=(epoch % visualize_every == 0 or epoch == 1))
+                if self.interrupt:
+                    return
                 val_losses.append(val_loss)
                 val_accuracies.append(val_acc)
                 print(f"Epoch {epoch} Train Loss: {train_loss:.4f} | Train Acc : {train_acc:.4f} | Val Loss : {val_loss:.4f} | Val Acc : {val_acc:.4f}")
+                if val_acc > self.best_val_acc:
+                    print(f"New best model found at epoch {epoch} (Val Acc: {val_acc:.4f})")
+                    self.best_val_acc = val_acc
+                    self.best_model_state = {k: v.clone() for k, v in self.model.state_dict().items()}
+
                 yield train_loss,train_acc,val_loss,val_acc,fig
             else:
                 print(f"Epoch {epoch} Train Loss: {train_loss:.4f} | Train Acc : {train_acc:.4f}")
                 yield train_loss,train_acc,None,None,None
+        if self.best_model_state is not None:
+            self.model.load_state_dict(self.best_model_state)
+            print(f"Best model (Val Acc: {self.best_val_acc:.4f}) loaded into trainer.model")
         yield train_losses,train_accuracies,val_losses,val_accuracies,None
 
     def validate(self,epoch, visualize=False):
@@ -177,6 +331,8 @@ class ModelTrainer:
         fig = None
         with torch.no_grad():
             for imgs, labels in val_pbar:
+                if self.interrupt:
+                    break
                 labels = labels.to(self.device)
 
                 outputs = self.model(imgs)