ajout dependances

TractionModel.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Jul  4 15:07:27 2021
+
+@author: AlexandreN
+"""
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+import torchvision
+
+
+class SingleTractionHead(nn.Module):
+
+    def __init__(self):
+        super(SingleTractionHead, self).__init__()
+        
+        self.head_locs = nn.Sequential(nn.Linear(2048, 1024),
+                                       nn.ReLU(),                 
+                                       nn.Dropout(p=0.3),
+                                       nn.Linear(1024, 4),
+                                       nn.Sigmoid()
+                                      )
+
+        # Head class should output the logits over the classe
+        self.head_class = nn.Sequential(nn.Linear(2048, 128),
+                                        nn.ReLU(),                          
+                                        nn.Dropout(p=0.3),
+                                        nn.Linear(128, 1))
+
+    def forward(self, features):
+        features = features.view(features.size()[0], -1)
+        
+        y_bbox  = self.head_locs(features)
+        y_class = self.head_class(features)
+
+        res = (y_bbox, y_class)
+        return res
+
+
+def create_model():
+    # setup the architecture of the model
+    feature_extractor = torchvision.models.resnet50(pretrained=True)
+    model_body = nn.Sequential(*list(feature_extractor.children())[:-1])
+    for param in model_body.parameters():
+        param.requires_grad = False
+    # Parameters of newly constructed modules have requires_grad=True by default
+    # num_ftrs = model_body.fc.in_features
+    
+    model_head = SingleTractionHead()
+    model = nn.Sequential(model_body, model_head)
+    return model
+
+
+def load_weights(model, path='model.pt'):
+    checkpoint = torch.load(path)
+    model.load_state_dict(checkpoint)
+    return model

app.py

@@ -5,7 +5,6 @@ import cv2
 import numpy as np
 import matplotlib.pyplot as plt
 
-from DataSet import QuestionDataSet
 import TractionModel as plup
 
 import random
@@ -40,7 +39,7 @@ vanilla_transform = torchvision.transforms.Compose([
                                                     torchvision.transforms.ToTensor(),
                                                     torchvision.transforms.Normalize(norm_mean, norm_std)])
 
-model = init_model("output/model/model-score0.96-f1_10.9-f1_20.99.pt")
+model = init_model("model-score0.96-f1_10.9-f1_20.99.pt")
 if torch.cuda.is_available():
     device = torch.device("cuda")
 else:

model-score0.96-f1_10.9-f1_20.99.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:acf13d3f6f4758fa68c8346478c9a7a5b1323cd96861a3c4266e7b8c438e4c18
+size 103808501