import torch.nn as nn
import torch
from torchvision.models import resnet50
import os
os.environ["TORCH_HOME"] = os.getcwd()

# reference : https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L166

class Resnet50Custom(nn.Module):
    def __init__(self,  object_type_output_dim=15):
        super().__init__()
        self.resnet50 = resnet = resnet50(weights="IMAGENET1K_V2")
        for param in self.resnet50.parameters():
            param.requires_grad = False
        self.resnet50.fc = nn.Linear(self.resnet50.fc.in_features, object_type_output_dim)
        self.ff_defect_classification = nn.Linear(self.resnet50.fc.in_features, 2)

    def forward(self, x):
        x = self.resnet50.conv1(x)
        x = self.resnet50.bn1(x)
        x = self.resnet50.relu(x)
        x = self.resnet50.maxpool(x)

        x = self.resnet50.layer1(x)
        x = self.resnet50.layer2(x)
        x = self.resnet50.layer3(x)
        x = self.resnet50.layer4(x)

        x = self.resnet50.avgpool(x)
        x = torch.flatten(x, 1)

        ff_object_classification_output = self.resnet50.fc(x)
        ff_defect_classification_output = self.ff_defect_classification(x)

        return ff_object_classification_output, ff_defect_classification_output