DataRaptor/ActionNet

1import torch2from torch import nn3from torchvision import transforms, models4 5class ActionClassifier(nn.Module):6    def __init__(self, train_last_nlayer, hidden_size, dropout, ntargets):7        super().__init__()8        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT, progress=True)9        modules = list(resnet.children())[:-1] # delete last layer10 11        self.resnet = nn.Sequential(*modules)12        for param in self.resnet[:-train_last_nlayer].parameters():13            param.requires_grad = False14            15        self.fc = nn.Sequential(16            nn.Flatten(),17            nn.BatchNorm1d(resnet.fc.in_features),18            nn.Dropout(dropout),19            nn.Linear(resnet.fc.in_features, hidden_size),20            nn.ReLU(),21            nn.BatchNorm1d(hidden_size),22            nn.Dropout(dropout),23            nn.Linear(hidden_size, ntargets),24            nn.Sigmoid()25        )26    27    def forward(self, x):28        x = self.resnet(x)29        x = self.fc(x)30        return x31 32 33def get_transform():34    transform = transforms.Compose([35        transforms.Resize([224, 244]),36        models.ResNet50_Weights.DEFAULT.transforms()37    ])38    return transform39 40# def get_transform():41#     transform = transforms.Compose([ 42#                 transforms.Resize([224, 244]), 43#                 transforms.ToTensor(),44#                 # std multiply by 255 to convert img of [0, 255]45#                 # to img of [0, 1]46#                 transforms.Normalize((0.485, 0.456, 0.406), 47#                                      (0.229*255, 0.224*255, 0.225*255))]48#             )49#     return transform50 51 52def get_model():53    model = ActionClassifier(0, 512, 0.2, 15)54    model.load_state_dict(torch.load('./model_weights.pth', map_location=torch.device('cpu')))55    return model56 57 58def get_class(index):59    ind2cat = [60        'calling',61        'clapping',62        'cycling',63        'dancing',64        'drinking',65        'eating',66        'fighting',67        'hugging',68        'laughing',69        'listening_to_music',70        'running',71        'sitting',72        'sleeping',73        'texting',74        'using_laptop'75    ]76    return ind2cat[index]77 78 79 80 81# img = Image.open('./inputs/Image_102.jpg').convert('RGB')82# #print(transform(img))83# img = transform(img)84# img = img.unsqueeze(dim=0)85# print(img.shape)86 87 88 89 90 91 92# model.eval()93# with torch.no_grad():94#     out = model(img)95#     out = nn.Softmax()(out).squeeze()96#     print(out.shape)97#     res = torch.argmax(out)98    99#     print(ind2cat[res])100 101 102 103 104