Spaces:
Runtime error
Runtime error
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision.transforms as transforms | |
| from torchvision import models | |
| import torchvision | |
| import torch | |
| import copy | |
| class ForkliftFrameClassifier_V0(nn.Module): | |
| def __init__(self, n_classes = 2, dropout = 0.15): | |
| super(ForkliftFrameClassifier_V0, self).__init__() | |
| self.dropout = dropout | |
| # N x 3 x 480 x 640 | |
| self.Conv1 = nn.Conv2d(3, 32, kernel_size=(8,8), stride=(3,5), padding=(3,1)) | |
| self.Bn1 = nn.BatchNorm2d(32) | |
| # N x 32 x 160 x 127 | |
| self.Conv2 = nn.Conv2d(32, 64, kernel_size=(8,8), stride=(5,5), padding=(0,0)) | |
| self.Bn2 = nn.BatchNorm2d(64) | |
| # N x 64 x 31 x 24 | |
| self.Maxpool1 = nn.MaxPool2d(kernel_size=(5,5), stride=(3,3), padding=(0,2)) | |
| # N x 64 x 9 x 8 | |
| self.Conv3 = nn.Conv2d(64, 64, kernel_size=(5,5), stride=(3,3), padding=(1,2)) | |
| self.Bn3 = nn.BatchNorm2d(64) | |
| # N x 64 x 3 x 3 | |
| self.Maxpool2 = nn.MaxPool2d(kernel_size=(3,3), stride=(1,1), padding=(0,0)) | |
| # N x 64 x 1 x 1 | |
| self.Linear1 = nn.Linear(64, 16) | |
| self.FC_out = nn.Linear(16, 1) if n_classes==2 else nn.Linear(64, n_classes) | |
| def forward(self, x): | |
| #print(x.shape) | |
| x = self.Conv1(x) | |
| x = self.Bn1(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Conv2(x) | |
| x = self.Bn2(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Maxpool1(x) | |
| #print(x.shape) | |
| x = self.Conv3(x) | |
| x = self.Bn3(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Maxpool2(x) | |
| #print(x.shape) | |
| x = x.reshape(x.shape[0], -1) | |
| #print(x.shape) | |
| x = self.Linear1(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.FC_out(x) | |
| return x | |
| class ForkliftFrameClassifier_V1(nn.Module): | |
| def __init__(self, n_classes = 2, dropout = 0.15): | |
| super(ForkliftFrameClassifier_V1, self).__init__() | |
| self.dropout = dropout | |
| # N x 3 x 240 x 240 | |
| self.Conv1 = nn.Conv2d(3, 32, kernel_size=(5,5), stride=(3,3), padding=(1,1)) | |
| self.Bn1 = nn.BatchNorm2d(32) | |
| # N x 32 x 80 x 80 | |
| self.Conv2 = nn.Conv2d(32, 64, kernel_size=(5,5), stride=(3,3), padding=(1,1)) | |
| self.Bn2 = nn.BatchNorm2d(64) | |
| # N x 64 x 26 x 26 | |
| self.Maxpool1 = nn.MaxPool2d(kernel_size=(5,5), stride=(3,3), padding=(1,1)) | |
| # N x 64 x 8 x 8 | |
| self.Conv3 = nn.Conv2d(64, 32, kernel_size=(3,3), stride=(2,2), padding=(1,1)) | |
| self.Bn3 = nn.BatchNorm2d(32) | |
| # N x 64 x 4 x 4 | |
| self.Maxpool2 = nn.MaxPool2d(kernel_size=(4,4), stride=(1,1), padding=(0,0)) | |
| # N x 64 x 1 x 1 | |
| #self.Linear1 = nn.Linear(64, 16) | |
| self.FC_out = nn.Linear(32, 1) if n_classes==2 else nn.Linear(64, n_classes) | |
| def forward(self, x): | |
| #print(x.shape) | |
| #print(x.shape) | |
| x = self.Conv1(x) | |
| x = self.Bn1(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Conv2(x) | |
| x = self.Bn2(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Maxpool1(x) | |
| #print(x.shape) | |
| x = self.Conv3(x) | |
| x = self.Bn3(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Maxpool2(x) | |
| #print(x.shape) | |
| x = x.reshape(x.shape[0], -1) | |
| #print(x.shape) | |
| #x = self.Linear1(x) | |
| #x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.FC_out(x) | |
| return x | |
| class ForkliftFrameClassifier_V2(nn.Module): | |
| def __init__(self, n_classes = 2, dropout = 0.15): | |
| super(ForkliftFrameClassifier_V2, self).__init__() | |
| self.dropout = dropout | |
| # N x 3 x 150 x 150 | |
| self.Conv1 = nn.Conv2d(3, 32, kernel_size=(5,5), stride=(3,3), padding=(1,1)) | |
| self.Bn1 = nn.BatchNorm2d(32) | |
| # N x 32 x 50 x 50 | |
| self.Conv2 = nn.Conv2d(32, 64, kernel_size=(5,5), stride=(3,3), padding=(1,1)) | |
| self.Bn2 = nn.BatchNorm2d(64) | |
| # N x 64 x 16 x 16 | |
| self.Maxpool1 = nn.MaxPool2d(kernel_size=(3,3), stride=(2,2), padding=(1,1)) | |
| # N x 64 x 8 x 8 | |
| self.Conv3 = nn.Conv2d(64, 32, kernel_size=(3,3), stride=(2,2), padding=(1,1)) | |
| self.Bn3 = nn.BatchNorm2d(32) | |
| # N x 64 x 4 x 4 | |
| self.Maxpool2 = nn.MaxPool2d(kernel_size=(4,4), stride=(1,1), padding=(0,0)) | |
| # N x 64 x 1 x 1 | |
| #self.Linear1 = nn.Linear(64, 16) | |
| self.FC_out = nn.Linear(32, 1) if n_classes==2 else nn.Linear(64, n_classes) | |
| def forward(self, x): | |
| #print(x.shape) | |
| #print(x.shape) | |
| x = self.Conv1(x) | |
| x = self.Bn1(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Conv2(x) | |
| x = self.Bn2(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Maxpool1(x) | |
| #print(x.shape) | |
| x = self.Conv3(x) | |
| x = self.Bn3(x) | |
| x = F.relu(x) | |
| x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.Maxpool2(x) | |
| #print(x.shape) | |
| x = x.reshape(x.shape[0], -1) | |
| #print(x.shape) | |
| #x = self.Linear1(x) | |
| #x = F.dropout(x,self.dropout) | |
| #print(x.shape) | |
| x = self.FC_out(x) | |
| return x | |
| class ForkliftFrameClassifier_PT1(nn.Module): | |
| def __init__(self, pretrained_model,n_out_last_layer, n_classes = 2): | |
| super(ForkliftFrameClassifier_PT1, self).__init__() | |
| self.pt_model = pretrained_model | |
| self.pt_model.classifier = nn.Linear(25088,1) if n_classes==2 else nn.Linear(n_out_last_layer, n_classes) | |
| for param in self.pt_model.classifier.parameters(): | |
| param.requires_grad = False | |
| def forward(self, x): | |
| x = self.pt_model(x) | |
| return x | |
| class CNN_Feature_Extractor(nn.Module): | |
| def __init__(self): | |
| super(CNN_Feature_Extractor, self).__init__() | |
| self.conv1 = nn.Conv2d(3, 10, kernel_size=(5,5), stride=(3,3)) | |
| self.conv2 = nn.Conv2d(10, 20, kernel_size=(5,5), stride=(2,2)) | |
| self.conv3 = nn.Conv2d(20, 30, kernel_size=(5,5), stride=(2,2)) | |
| def forward(self, i): | |
| x = i.view(-1, i.shape[2], i.shape[3], i.shape[4]) | |
| x = F.relu(self.conv1(x)) | |
| x = F.relu(self.conv2(x)) | |
| x = F.relu(self.conv3(x)) | |
| x = nn.AvgPool2d(3)(x) | |
| x = x.view(i.shape[0], i.shape[1], -1) | |
| return x | |
| class LSTM(nn.Module): | |
| def __init__(self, seq_len, hidden_size,out_size): | |
| super(LSTM, self).__init__() | |
| self.lstm = nn.LSTM(750, hidden_size) | |
| self.fc = nn.Linear(seq_len*hidden_size, out_size) | |
| def forward(self, x): | |
| x, _ = self.lstm(x) | |
| x = x.view(x.shape[0], -1) | |
| x = self.fc(x) | |
| return x | |
| class Full_LSTM(nn.Module): | |
| def __init__(self,seq_len = 15, hidden_size = 100, out_size = 512): | |
| super(Full_LSTM, self).__init__() | |
| self.net_cnn = CNN_Feature_Extractor() | |
| self.net_lstm = LSTM(seq_len, hidden_size, out_size) | |
| self.classifier = nn.Sequential(nn.Linear(out_size, 16), | |
| nn.Dropout(0.3), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| def forward(self, x): | |
| # x.size() == (B,L,C,H,W) | |
| # B : Batch size | |
| # L : Sequence Length = 15 | |
| # C : Channels = 3 | |
| # H : Heigth = 224 | |
| # W : Width = 224 | |
| x = self.net_cnn(x) | |
| x = self.net_lstm(x) | |
| x = self.classifier(x) | |
| return x | |
| class Full_CNN(nn.Module): | |
| def __init__(self): | |
| super(Full_CNN, self).__init__() | |
| self.model = torchvision.models.resnet18(pretrained=True) | |
| #for param in model.parameters(): | |
| # param.requires_grad = False | |
| self.model.fc = nn.Sequential(nn.Linear(512, 16), | |
| nn.Dropout(0.3), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| def forward(self, x): | |
| # x.size() == (B,L,C,H,W) | |
| # B : Batch size | |
| # L : Sequence Length = 15 | |
| # C : Channels = 3 | |
| # H : Heigth = 224 | |
| # W : Width = 224 | |
| x = self.model(x[:,0,:]) | |
| return x | |
| class Full_Model(nn.Module): | |
| def __init__(self,seq_len = 15, hidden_size = 100, classifier_size = 512, dropout = 0.4, cnn_model_path = None, lstm_model_path= None): | |
| super(Full_Model, self).__init__() | |
| self.CNN_Part = Full_CNN() | |
| if cnn_model_path is not None: | |
| self.CNN_Part.model.load_state_dict(torch.load(cnn_model_path)['model_state_dict']) | |
| self.CNN_classifier = copy.deepcopy(self.CNN_Part.model.fc) | |
| else: | |
| self.CNN_classifier = nn.Sequential(nn.Linear(classifier_size, 16), | |
| nn.Dropout(dropout), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| self.CNN_Part.model.fc = nn.Sequential(nn.Linear(classifier_size, classifier_size), | |
| nn.Dropout(dropout), | |
| nn.ReLU()) | |
| self.LSTM_Part = Full_LSTM(seq_len, hidden_size, classifier_size) | |
| if lstm_model_path is not None: | |
| self.LSTM_Part.load_state_dict(torch.load(lstm_model_path)['model_state_dict']) | |
| self.LSTM_classifier = copy.deepcopy(self.LSTM_Part.classifier) | |
| else: | |
| self.LSTM_classifier = nn.Sequential(nn.Linear(classifier_size, 16), | |
| nn.Dropout(dropout), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| self.Finalclassifier = nn.Sequential(nn.Linear(classifier_size, 16), | |
| nn.Dropout(dropout), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| def forward(self, x): | |
| # x.size() == (B,L,C,H,W) | |
| # B : Batch size | |
| # L : Sequence Length = 15 | |
| # C : Channels = 3 | |
| # H : Heigth = 224 | |
| # W : Width = 224 | |
| cnn_out = self.CNN_Part(x) | |
| # xcnn : (B,512) | |
| lstm_out = self.LSTM_Part.net_cnn(x) | |
| lstm_out = self.LSTM_Part.net_lstm(lstm_out) | |
| # xlstm : (B,512) | |
| out = cnn_out + lstm_out | |
| cnn_out = self.CNN_classifier(cnn_out) | |
| lstm_out = self.LSTM_classifier(lstm_out) | |
| out = self.Finalclassifier(out) | |
| return (cnn_out, lstm_out, out) | |
| # Se establece que el modelo final consista en cargar un modelo Full, pero considerando solamente la salida y el | |
| # forward de la componente LSTM. Esto se hace así (Cargar inclusive los pesos de LSTM) por si en el futuro se decidiera que estos pesos pueden ser útiles. | |
| class Final_CNN_Model(nn.Module): | |
| def __init__(self,seq_len = 15, hidden_size = 100, classifier_size = 512, dropout = 0.4, cnn_model_path = None, lstm_model_path= None): | |
| super(Final_CNN_Model, self).__init__() | |
| self.CNN_Part = Full_CNN() | |
| if cnn_model_path is not None: | |
| self.CNN_Part.model.load_state_dict(torch.load(cnn_model_path)['model_state_dict']) | |
| self.CNN_classifier = copy.deepcopy(self.CNN_Part.model.fc) | |
| else: | |
| self.CNN_classifier = nn.Sequential(nn.Linear(classifier_size, 16), | |
| nn.Dropout(dropout), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| self.CNN_Part.model.fc = nn.Sequential(nn.Linear(classifier_size, classifier_size), | |
| nn.Dropout(dropout), | |
| nn.ReLU()) | |
| self.LSTM_Part = Full_LSTM(seq_len, hidden_size, classifier_size) | |
| if lstm_model_path is not None: | |
| self.LSTM_Part.load_state_dict(torch.load(lstm_model_path)['model_state_dict']) | |
| self.LSTM_classifier = copy.deepcopy(self.LSTM_Part.classifier) | |
| else: | |
| self.LSTM_classifier = nn.Sequential(nn.Linear(classifier_size, 16), | |
| nn.Dropout(dropout), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| self.Finalclassifier = nn.Sequential(nn.Linear(classifier_size, 16), | |
| nn.Dropout(dropout), | |
| nn.ReLU(), | |
| nn.Linear(16,1)) | |
| def forward(self, x): | |
| # x.size() == (B,L,C,H,W) | |
| # B : Batch size | |
| # L : Sequence Length = 15 | |
| # C : Channels = 3 | |
| # H : Heigth = 224 | |
| # W : Width = 224 | |
| cnn_out = self.CNN_Part(x) | |
| cnn_out = self.CNN_classifier(cnn_out) | |
| return cnn_out | |
| # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| # torch.cuda.empty_cache() | |
| # x1 = torch.rand((64, 15, 3, 224 , 224)) | |
| # model = Full_Model(cnn_model_path = 'Best_model_4.pt', lstm_model_path= 'Best_model_10.pt') | |
| # model.to(device) | |
| # x1 = x1.to(device) | |
| # out1 = model(x1) | |