extra

app.py

@@ -33,7 +33,7 @@ model = YOLOv3Lightning()
 model.load_state_dict(torch.load("yolov3_608_ckpt_40.pth", map_location=torch.device('cpu')), strict=False)
 model.setup(stage="test")
 
-IMAGE_SIZE = 416
+IMAGE_SIZE = 608
 transforms = A.Compose(
     [
         A.LongestMaxSize(max_size=IMAGE_SIZE),

config.py

@@ -10,7 +10,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 # seed_everything()  # If you want deterministic behavior
 NUM_WORKERS = 4
 BATCH_SIZE = 16
-IMAGE_SIZE = 416
+IMAGE_SIZE = 608
 NUM_CLASSES = 1
 LEARNING_RATE = 1e-5
 WEIGHT_DECAY = 1e-4

main_yolov3_lightening.py

@@ -22,7 +22,7 @@ class YOLOv3Lightning(pl.LightningModule):
         super().__init__()
         self.config = config
         #self.model = YOLOv3(num_classes=self.config.NUM_CLASSES)
-        self.model = load_model("yolov3_divots_608.cfg", "yolov3_608_ckpt_40.pth")
+        self.model = load_model("yolov3_divots_608.cfg")
         self.loss_fn = YoloLoss()
         if lr_value == 0:
           self.learning_rate = self.config.LEARNING_RATE

pytorchyolo/models.py

@@ -316,7 +316,22 @@ class Darknet(nn.Module):
 
         fp.close()
 
+def load_model(model_path):
+    """Loads the yolo model from file.
 
+    :param model_path: Path to model definition file (.cfg)
+    :type model_path: str
+    :param weights_path: Path to weights or checkpoint file (.weights or .pth)
+    :type weights_path: str
+    :return: Returns model
+    :rtype: Darknet
+    """
+    device = torch.device("cuda" if torch.cuda.is_available()
+                          else "cpu")  # Select device for inference
+    model = Darknet(model_path).to(device)
+    
+    return model
+    
 def load_model(model_path, weights_path=None):
     """Loads the yolo model from file.