import torch 
import torch.nn  as nn
from bn import batch_norm
from cbam import CBAM
class residual(nn.Module):
    def __init__(self, inp, out, stride=1):
        super().__init__()
        self.bn1 = batch_norm(inp)
        self.conv1 = nn.Conv2d(inp, out, kernel_size=3, padding=1, stride=stride)
        self.bn2 = batch_norm(out)
        self.conv2 = nn.Conv2d(out, out, kernel_size=3, padding=1, stride=1)
        # skip connection
        self.concat = nn.Conv2d(inp, out, kernel_size=1, padding=0, stride=stride)
        # Add CBAM
        self.cbam = CBAM(out)

    def forward(self, input):
        x = self.bn1(input)
        x = self.conv1(x)
        x = self.bn2(x)
        x = self.conv2(x)
        x = self.cbam(x)  # Apply CBAM
        skip = self.concat(input)
        skip = x + skip
        return skip