200 lines
7.2 KiB
Python
200 lines
7.2 KiB
Python
# Dataset utils and dataloaders
|
||
|
||
import glob
|
||
import logging
|
||
import math
|
||
import os
|
||
import random
|
||
import shutil
|
||
import time
|
||
from itertools import repeat
|
||
from multiprocessing.pool import ThreadPool
|
||
from pathlib import Path
|
||
from threading import Thread
|
||
import re
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from PIL import Image, ExifTags
|
||
from torch.utils.data import Dataset
|
||
from tqdm import tqdm
|
||
|
||
# Parameters
|
||
help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
|
||
img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng'] # acceptable image suffixes
|
||
vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Get orientation exif tag
|
||
for orientation in ExifTags.TAGS.keys():
|
||
if ExifTags.TAGS[orientation] == 'Orientation':
|
||
break
|
||
|
||
|
||
|
||
|
||
class LoadImages: # for inference
|
||
def __init__(self, path, stride=32, img_size=640):
|
||
p = str(Path(path)) # os-agnostic
|
||
p = os.path.abspath(p) # absolute path
|
||
if '*' in p:
|
||
files = sorted(glob.glob(p, recursive=True)) # glob
|
||
elif os.path.isdir(p):
|
||
files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
|
||
elif os.path.isfile(p):
|
||
files = [p] # files
|
||
else:
|
||
raise Exception(f'ERROR: {p} does not exist')
|
||
|
||
images = [x for x in files if x.split('.')[-1].lower() in img_formats]
|
||
videos = [x for x in files if x.split('.')[-1].lower() in vid_formats]
|
||
ni, nv = len(images), len(videos)
|
||
|
||
self.img_size = img_size
|
||
self.files = images + videos
|
||
self.nf = ni + nv # number of files
|
||
self.video_flag = [False] * ni + [True] * nv
|
||
self.mode = 'image'
|
||
self.stride = stride
|
||
if any(videos):
|
||
self.new_video(videos[0]) # new video
|
||
else:
|
||
self.cap = None
|
||
assert self.nf > 0, f'No images or videos found in {p}. ' \
|
||
f'Supported formats are:\nimages: {img_formats}\nvideos: {vid_formats}'
|
||
|
||
def __iter__(self):
|
||
self.count = 0
|
||
return self
|
||
|
||
def __next__(self):
|
||
if self.count == self.nf:
|
||
raise StopIteration
|
||
path = self.files[self.count]
|
||
|
||
if self.video_flag[self.count]:
|
||
# Read video
|
||
self.mode = 'video'
|
||
ret_val, img0 = self.cap.read() # ret：表示读取是否成功的布尔值；
|
||
if not ret_val:
|
||
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
||
_, img0 = self.cap.read()
|
||
# print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='\n')
|
||
|
||
else:
|
||
# Read image
|
||
self.count += 1
|
||
img0 = cv2.imread(path) # BGR
|
||
assert img0 is not None, 'Image Not Found ' + path
|
||
# print(f'image {self.count}/{self.nf} {path}: ', end='\n')
|
||
|
||
|
||
img0 = letterbox(img0, new_shape=self.img_size, stride =self.stride)[0]
|
||
# img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
|
||
return img0
|
||
|
||
def new_video(self, path):
|
||
self.frame = 0
|
||
self.cap = cv2.VideoCapture(path)
|
||
self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
|
||
def __len__(self):
|
||
return self.nf # number of files
|
||
|
||
|
||
class LoadStreams: # multiple IP or RTSP cameras
|
||
def __init__(self, sources='streams.txt', img_size=640):
|
||
self.mode = 'stream'
|
||
self.img_size = img_size
|
||
|
||
if os.path.isfile(sources):
|
||
with open(sources, 'r') as f:
|
||
sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]
|
||
else:
|
||
sources = [sources]
|
||
|
||
n = len(sources)
|
||
self.imgs = [None] * n
|
||
self.sources = [clean_str(x) for x in sources] # clean source names for later
|
||
for i, s in enumerate(sources):
|
||
# Start the thread to read frames from the video stream
|
||
print(f'{i + 1}/{n}: {s}... ', end='')
|
||
cap = cv2.VideoCapture(eval(s) if s.isnumeric() else s)
|
||
assert cap.isOpened(), f'Failed to open {s}'
|
||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
fps = cap.get(cv2.CAP_PROP_FPS) % 100
|
||
_, self.imgs[i] = cap.read() # guarantee first frame
|
||
thread = Thread(target=self.update, args=([i, cap]), daemon=True)
|
||
print(f' success ({w}x{h} at {fps:.2f} FPS).')
|
||
thread.start()
|
||
print('') # newline
|
||
|
||
# check for common shapes
|
||
s = np.stack([letterbox(x, new_shape=self.img_size)[0].shape for x in self.imgs], 0) # inference shapes
|
||
self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
|
||
if not self.rect:
|
||
print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
|
||
|
||
def update(self, index, cap):
|
||
# Read next stream frame in a daemon thread
|
||
n = 0
|
||
while cap.isOpened():
|
||
n += 1
|
||
# _, self.imgs[index] = cap.read()
|
||
cap.grab()
|
||
if n == 4: # read every 4th frame
|
||
_, self.imgs[index] = cap.retrieve()
|
||
n = 0
|
||
time.sleep(0.01) # wait time
|
||
|
||
def __iter__(self):
|
||
self.count = -1
|
||
return self
|
||
|
||
def __next__(self):
|
||
self.count += 1
|
||
img0 = self.imgs.copy()
|
||
|
||
return img0
|
||
|
||
def __len__(self):
|
||
return 0 # 1E12 frames = 32 streams at 30 FPS for 30 years
|
||
|
||
|
||
def clean_str(s):
|
||
# Cleans a string by replacing special characters with underscore _
|
||
return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
|
||
|
||
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
|
||
# Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
|
||
shape = img.shape[:2] # current shape [height, width]
|
||
if isinstance(new_shape, int):
|
||
new_shape = (new_shape, new_shape)
|
||
|
||
# Scale ratio (new / old)
|
||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||
if not scaleup: # only scale down, do not scale up (for better test mAP)
|
||
r = min(r, 1.0)
|
||
|
||
# Compute padding
|
||
ratio = r, r # width, height ratios
|
||
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
||
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
||
if auto: # minimum rectangle
|
||
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
|
||
elif scaleFill: # stretch
|
||
dw, dh = 0.0, 0.0
|
||
new_unpad = (new_shape[1], new_shape[0])
|
||
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
|
||
|
||
dw /= 2 # divide padding into 2 sides
|
||
dh /= 2
|
||
|
||
if shape[::-1] != new_unpad: # resize
|
||
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
||
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
|
||
return img, ratio, (dw, dh) |