Update bgremover.py

bgremover.py

@@ -351,20 +351,22 @@ class DamageClassifier():
             model_filepath = r"\\CATALOGUE.CGIARAD.ORG\AcceleratedBreedingInitiative\1.Data\16. Spidermites_AdrianK\best_models\short_resnet18_SpidermitesModel.pth"
             model_filepath = model_filepath = os.path.join(MODEL_PATH, "short_resnet18_SpidermitesModel.pth")
             model = models.resnet18(weights='IMAGENET1K_V1')
-        
-        #Add fully connected layer at the end with num_classes as output
-        num_ftrs = model.fc.in_features
-        model.fc = nn.Linear(num_ftrs, 4)
 
-        if torch.cuda.is_available():
-            model.load_state_dict(torch.load(model_filepath))
-            model.cuda()
-        else:
-            model.load_state_dict(torch.load(model_filepath, map_location='cpu'))
-        model.eval()
-
-        self.model = model
-        self.model_name = model_name
+        if not model is None:
+        
+            #Add fully connected layer at the end with num_classes as output
+            num_ftrs = model.fc.in_features
+            model.fc = nn.Linear(num_ftrs, 4)
+    
+            if torch.cuda.is_available():
+                model.load_state_dict(torch.load(model_filepath))
+                model.cuda()
+            else:
+                model.load_state_dict(torch.load(model_filepath, map_location='cpu'))
+            model.eval()
+    
+            self.model = model
+            self.model_name = model_name
 
         return