export tensorrt engine and add [AnomalyDINO: Boosting Patch-based Few-shot Anomaly Detection with DINOv2] method

[lalala0814]

export tensorrt

    #src\anomalib\engine\engine.py 825
    elif export_type == ExportType.TENSORRT:
        # NOTE: TensorRT export is only supported for batch_size=1 or dynamic shapes.
        exported_model_path=model.to_engine(
            export_root=export_root,
            model_file_name=model_file_name,
            input_size=input_size,
            #batch_size=batch_size,
        )
    #src\anomalib\models\components\base\export_mixin.py
    def to_engine(
        self,
        export_root: Path | str,
        model_file_name: str = "model",
        input_size: tuple[int, int] | None = None,
        batch_size:  int = 1,
        compatible_mode: bool =True,
    ):
        import tensorrt as trt
        logger = trt.Logger(trt.Logger.INFO)
        builder = trt.Builder(logger)
        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        parser = trt.OnnxParser(network, logger)
        onnx_path=self.to_onnx(export_root, model_file_name, input_size,batch_size)
        export_root= _create_export_root(export_root, ExportType.TENSORRT)
        engine_export_path=export_root / (model_file_name + ".engine")
        with open(onnx_path, "rb") as f:
            parser.parse(f.read())


        config = builder.create_builder_config()

        config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
        config.max_workspace_size = 1 << 30  
        is_trt10 = int(trt.__version__.split(".")[0]) >= 10 
        if compatible_mode:
            if is_trt10:
                config.set_hardware_compatibility_level(trt.HardwareCompatibilityLevel.AMPERE_PLUS)
            else:
                if hasattr(config, 'hardware_compatibility_level'):
                    config.hardware_compatibility_level = trt.HardwareCompatibilityLevel.AMPERE_PLUS
                else:
                    self.g_logger.warning("hardware_compatibility_level not supported in this TensorRT version")
        input_tensor = network.get_input(0)
        input_shape = input_tensor.shape  
        if input_shape[0] == -1:
            #input_shape[0] = 1
            min_shape = (1, *input_shape[1:])
            opt_shape = (1, *input_shape[1:])
            max_shape = (50, *input_shape[1:])
            profile = builder.create_optimization_profile()
            profile.set_shape(input_tensor.name, min=min_shape, opt=opt_shape, max=max_shape)
            config.add_optimization_profile(profile)

        all_outputs = [network.get_output(i) for i in range(network.num_outputs)]
        for out in all_outputs:
            if out.name not in  ['anomaly_map']:
                network.unmark_output(out)

        engine = builder.build_engine(network, config)
        with open(engine_export_path, "wb") as f:
            f.write(engine.serialize())

anomaly_dinov2

# src\anomalib\models\image\anomaly_dinov2\lightning_model.py
class AnomalyDinoV2(MemoryBankMixin,AnomalibModule):
    def __init__(
        self,
        backbone: DinoModelType | str = DinoModelType.DINOV2_S,
        k_shot: int = 1,
        knn_neighbors: int = 1,
        img_size: tuple[int, int] = (256, 256),
        pre_processor: nn.Module | bool = True,
        post_processor: nn.Module | bool = True,
        evaluator: Evaluator | bool = True,
        visualizer: Visualizer | bool = True,
        use_faiss:bool=True,
    ) -> None:
        super().__init__(
            pre_processor=pre_processor,
            post_processor=post_processor,
            evaluator=evaluator,
            visualizer=visualizer,
        )
        self.img_size=img_size
        self.k_shot=k_shot
        self.model:AnomalyDinoModel=AnomalyDinoModel(backbone,k_shot,knn_neighbors,img_size,use_faiss)

    @staticmethod
    def configure_optimizers() -> None:
        return

    def training_step(self, batch: Batch, *args, **kwargs) -> None:
        del args, kwargs  # These variables are not used.
        feature = self.model(batch.image)
        return torch.tensor(0.0, requires_grad=True, device=self.device)

    @property
    def trainer_arguments(self) -> dict[str, int | float]:
        return {"max_epochs": 1, "val_check_interval": 1.0, "num_sanity_val_steps": 0}

#src\anomalib\models\image\anomaly_dinov2\torch_model.py
class AnomalyDinoModel(DynamicBufferMixin,nn.Module):
    def __init__(
        self,
        model_type:DinoModelType|str=DinoModelType.DINOV2_S,
       
        k_shots:int=1,
        knn_neighbors:int=1,
        img_size:tuple[int,int]=(256,256),
        use_faiss:bool=True,
        half_precision:bool=False,
    ):
        super().__init__()
        if model_type==DinoModelType.DINOV2_S:
            self.model=dinov2_vits14()
        self.model.eval()
        self.half_precision=half_precision
        self.k_shots=k_shots
        self.features_ref:torch.Tensor
        self.register_buffer("features_ref", torch.empty(0))
        self.img_size=img_size
        self.cropped_height,self.cropped_width = img_size[0] - img_size[0] % self.model.patch_size, img_size[1] - img_size[1] % self.model.patch_size
        self.grid_size=(self.cropped_height // self.model.patch_size, self.cropped_width // self.model.patch_size)
        self.anomaly_map_generator:KnnDistance=KnnDistance(use_faiss=use_faiss,knn_neighbors=knn_neighbors)

    def extract_features(self, image_tensor:torch.Tensor):
        image_tensor=image_tensor[:,:, :self.cropped_height, :self.cropped_width]
        with torch.inference_mode():
            if self.half_precision:
                image_batch = image_tensor.half()
            else:
                image_batch = image_tensor

            tokens = self.model.get_intermediate_layers(image_batch)[0]
            if (self.features_ref is not None and 
                self.k_shots != -1 and 
                self.features_ref.shape[0] >= self.k_shots):
                return tokens
            if self.features_ref is None:
                if self.k_shots == -1:  
                    self.features_ref = tokens.clone()
                else: 
                    self.features_ref = tokens[:min(tokens.shape[0], self.k_shots)].clone()
                return tokens
            if self.training:
                if self.k_shots==-1:
                    self.features_ref=torch.cat([self.features_ref, tokens], dim=0)
                else:
                    if (self.features_ref.shape[0]+tokens.shape[0])<self.k_shots:
                        self.features_ref=torch.cat([self.features_ref, tokens], dim=0)
                    else:
                        self.features_ref=torch.cat([self.features_ref, tokens], dim=0)[:self.k_shots]
        return tokens
    

    
    def forward(self,input_tensor:torch.Tensor):
        tokens=self.extract_features(input_tensor)
        pred_map=self.anomaly_map_generator.fit(self.features_ref,tokens,self.grid_size,self.img_size)
        pred_scores=torch.amax(pred_map, dim=(-2, -1))
        return InferenceBatch(pred_score=pred_scores,anomaly_map=pred_map)
        
#src\anomalib\models\image\anomaly_dinov2\anomaly_map.py
import torch
import torch.nn as nn
import torch.nn.functional as F
def mean_top1p( distances: torch.Tensor) -> torch.Tensor:
    top1 = distances[:, 0]  
    N = top1.size(0)
    if int(N * 0.01) == 0:
        return top1.max()
    else:
        k = int(N * 0.01)
        return torch.topk(top1, k=k, largest=True).values.mean()

class KnnDistance(nn.Module):  
    def __init__(
        self,
        knn_neighbors: int = 1,
        img_size: tuple[int, int] = (256, 256),
        use_faiss:bool=False,
        use_chunk: bool = False
    ) -> None:
        super().__init__()
        self.knn_neighbors = knn_neighbors
        self.img_size = img_size
        self.use_chunk = use_chunk
        self.chunk_size = 2
        #self.register_buffer("features", torch.empty(0))  

    def _normalize(self, x: torch.Tensor) -> torch.Tensor:
        return F.normalize(x, p=2, dim=1)

    def _compute_cosine_knn(self, db: torch.Tensor, query: torch.Tensor, k: int):
        db = self._normalize(db)
        query = self._normalize(query)
        similarity = torch.matmul(query, db.T)  # [N_query, N_db]
        distances = 1 - similarity  
        topk_vals, _ = torch.topk(distances, k, largest=False, dim=1)
        return topk_vals 
    


    @torch.no_grad()
    def forward(self, features: torch.Tensor, embedding: torch.Tensor, grid_size,img_size:tuple[int,int]):
        outputs=[]
        
        B, _, D = embedding.shape
        db = features.reshape(-1, features.shape[-1]) 
        for b in range(B):
            query = embedding[b].reshape(-1, D)    # [N_query, D]
            distances = self._compute_cosine_knn(db, query, self.knn_neighbors)  # [N_query, k]
            
            patch_distances = distances[:, 0].reshape(grid_size)  
            patch_distances = patch_distances.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
            
            resized = F.interpolate(
                patch_distances,
                size=img_size,
                mode="bilinear",
                align_corners=False
            )  
            outputs.append(resized)


        output_tensor = torch.cat(outputs, dim=0) 
        torch.cuda.empty_cache()

        return  output_tensor


    def fit(self, features: torch.Tensor, embedding: torch.Tensor, grid_size,img_size:tuple[int,int]):
        return self.forward(features, embedding, grid_size,img_size)

# train
def train():
    input_size = (512, 512)
    train_transform = Compose(
        [
            Resize(input_size, interpolation=InterpolationMode.BICUBIC, antialias=True),
            ToTensor(),
        ],
    )
    eval_transform = Compose(
        [
            Resize(input_size, interpolation=InterpolationMode.BICUBIC,antialias=True),
            ToTensor(),
        ],
    )
   # Create the datamodule
    datamodule = Folder(
        name="xx",
        root="xx",
        normal_dir="xx",
        num_workers=1,
        val_split_mode='same_as_test',
        train_augmentations=train_transform,
        val_augmentations=eval_transform,
        train_batch_size=1,
        eval_batch_size=1,
        test_split_ratio=0.01,  
    )
    datamodule.prepare_data()
    datamodule.setup()
    pre_process=PreProcessor(train_transform)
    model = AnomalyDinoV2(pre_processor=pre_process,k_shot=-1,img_size=input_size)
    engine = Engine(max_epochs=100)
    engine.fit(datamodule=datamodule, model=model)
    engine.export(model=model, export_type='onnx', input_size=input_size,batch_size=1)
    engine.export(model=model, export_type='openvino', input_size=input_size,batch_size=1)
    engine.export(model=model, export_type='engine', input_size=input_size,batch_size=1)