diff --git a/detect.py b/detect.py index da63114f..b71d5d9d 100755 --- a/detect.py +++ b/detect.py @@ -37,7 +37,7 @@ def detect( model.to(device).eval() # Set Dataloader - dataloader = load_images(images, img_size=img_size) + dataloader = LoadImages(images, img_size=img_size) # Get classes and colors classes = load_classes(parse_data_cfg('cfg/coco.data')['names']) diff --git a/test.py b/test.py index 69c21db1..e0d5deb3 100644 --- a/test.py +++ b/test.py @@ -36,8 +36,8 @@ def test( model.to(device).eval() # Get dataloader - # dataloader = torch.utils.data.DataLoader(load_images_with_labels(test_path), batch_size=batch_size) # pytorch - dataloader = load_images_and_labels(test_path, batch_size=batch_size, img_size=img_size) + # dataloader = torch.utils.data.DataLoader(LoadImagesAndLabels(test_path), batch_size=batch_size) # pytorch + dataloader = LoadImagesAndLabels(test_path, batch_size=batch_size, img_size=img_size) mean_mAP, mean_R, mean_P = 0.0, 0.0, 0.0 print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP')) diff --git a/train.py b/train.py index 383b5c6c..e7290b13 100644 --- a/train.py +++ b/train.py @@ -43,7 +43,7 @@ def train( model = Darknet(cfg, img_size) # Get dataloader - dataloader = load_images_and_labels(train_path, batch_size, img_size, multi_scale=multi_scale, augment=True) + dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, multi_scale=multi_scale, augment=True) lr0 = 0.001 if resume: diff --git a/utils/datasets.py b/utils/datasets.py index 95e86fcd..ef64c4e9 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -12,7 +12,7 @@ import torch from utils.utils import xyxy2xywh -class load_images(): # for inference +class LoadImages: # for inference def __init__(self, path, img_size=416): if os.path.isdir(path): image_format = ['.jpg', '.jpeg', '.png', '.tif'] @@ -55,7 +55,7 @@ class load_images(): # for inference return self.nF # number of files -class load_images_and_labels(): # for training +class LoadImagesAndLabels: # for training def __init__(self, path, batch_size=1, img_size=608, multi_scale=False, augment=False): self.path = path with open(path, 'r') as file: