extra

app.py

@@ -30,7 +30,7 @@ import gradio as gr
 import os
 
 model = YOLOv3Lightning()
-model.load_state_dict(torch.load("yolov3_model.pth", map_location=torch.device('cpu')), strict=False)
+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

config.py

@@ -11,7 +11,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 NUM_WORKERS = 4
 BATCH_SIZE = 16
 IMAGE_SIZE = 416
-NUM_CLASSES = 20
+NUM_CLASSES = 1
 LEARNING_RATE = 1e-5
 WEIGHT_DECAY = 1e-4
 NUM_EPOCHS = 100

main_yolov3_lightening.py

@@ -10,6 +10,13 @@ import config
 from torch.optim.lr_scheduler import OneCycleLR
 
 
+from pytorchyolo.models import load_model
+from pytorchyolo.utils.utils import load_classes, ap_per_class, get_batch_statistics, non_max_suppression, to_cpu, xywh2xyxy, print_environment_info
+from pytorchyolo.utils.datasets import ListDataset
+from pytorchyolo.utils.transforms import DEFAULT_TRANSFORMS
+from pytorchyolo.utils.parse_config import parse_data_config
+
+
 class YOLOv3Lightning(pl.LightningModule):
     def __init__(self,lr_value=0):
         super().__init__()