diff --git a/yolov5/detect.py b/yolov5/detect.py index e37abd220dbdc96b9ac12d403c25eb28ee3436a0..6f18cea9e97967aad15941a100a0aac13cda98dc 100755 --- a/yolov5/detect.py +++ b/yolov5/detect.py @@ -252,7 +252,7 @@ def run( def parse_opt(): parser = argparse.ArgumentParser() parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL') - parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)') + parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam) or a .csv with path column') parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path') parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold') diff --git a/yolov5/utils/dataloaders.py b/yolov5/utils/dataloaders.py index 3fbdab5138188e34036093f83c567636d4ff5916..e8eee11b2fae4cf73a31df7006ec88fc27886f2e 100755 --- a/yolov5/utils/dataloaders.py +++ b/yolov5/utils/dataloaders.py @@ -241,19 +241,25 @@ class LoadScreenshots: class LoadSpectros: def __init__(self, folder, sampleDur, rf, window, hop, low, high, cmap, minimum, img_size, stride=32, auto=True): self.folder, self.sampleDur, self.rf, self.window, self.hop, self.low, self.high, self.cmap, self.minimum, self.img_size, self.stride, self.auto = folder, sampleDur, rf, window, hop, low, high, cmap, minimum, img_size, stride, auto - self.files = os.listdir(folder) + self.mode = 'image' self.samples = [] - for fn in tqdm(os.listdir(folder), desc='Dataset initialization', leave=False): + if folder.split('.')[-1] == 'csv': + df_files = pd.read_csv(folder) + files = df_files.path.tolist() + else: + files = os.listdir(folder) + files = [os.path.join(folder, fn) for fn in files] + for fn in tqdm(files, desc='Dataset initialization', leave=False): try: - info = sf.info(os.path.join(folder, fn)) + info = sf.info(os.path.join(fn)) duration, fs = info.duration, info.samplerate self.samples.extend([{'fn':fn, 'offset':offset, 'fs':fs} for offset in np.arange(0, duration+.01 - sampleDur, sampleDur)]) except: print(f'failed with {fn}') continue self.nf = len(self.samples) - + def __len__(self): return len(self.samples)