rajora/binary-classification-galaxies

1import os2import torch3import torch.nn as nn4from efficientnet_pytorch import EfficientNet5 6class EffNet(nn.Module):7    def __init__(self, n_classes):8        super(EffNet, self).__init__()9        self.b4 = EfficientNet.from_pretrained('efficientnet-b0')10        self.drop = nn.Dropout(0.2)11        self.fc = nn.Linear(1000, n_classes)12 13    def forward(self, image):14        x = self.b4(image)15        x = self.drop(x)16        out = self.fc(x)17        return out18 19def load_model():20    device = torch.device("cpu")21    net = EffNet(n_classes=2).to(device)22    model_path = os.path.join(os.path.dirname(__file__), 'models', 'modelo_galaxias.pth')23    net.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))  # Adjust path if needed24    net.eval()25    return net