from torch import nn
import torch
from cassle.model.cbam_resnet import resnet18_cbam
from torchvision.models import resnet18

class torchLinearModel(nn.Module):
    def __init__(self, pretrained_path, num_classes):
        super().__init__()

        # self.backbone = resnet18_cbam()
        # self.backbone.inplanes = 512
        # self.backbone.linear = nn.Identity()
        self.backbone = resnet18()
        self.backbone.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=2, bias=False)
        self.backbone.maxpool = nn.Identity()
        self.backbone.fc = nn.Identity()
        self.classifier = nn.Linear(self.backbone.inplanes, num_classes)

        state_dict = torch.load(pretrained_path)['state_dict']
        self.load_state_dict(state_dict, strict=True)
        print(f"Loaded {pretrained_path}")

    def forward(self, x):
        x = self.backbone(x)
        x = self.classifier(x)
        return x
