# 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 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
from utils.general import check_requirements, xyxy2xywh, xywh2xyxy, xywhn2xyxy, xyn2xy, segment2box, segments2boxes, \
resample_segments, clean_str
from utils.torch_utils import torch_distributed_zero_first
# Parameters
help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # 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
def get_hash(files):
# Returns a single hash value of a list of files
return sum(os.path.getsize(f) for f in files if os.path.isfile(f))
def exif_size(img):
# Returns exif-corrected PIL size
s = img.size # (width, height)
try:
rotation = dict(img._getexif().items())[orientation]
if rotation == 6: # rotation 270
s = (s[1], s[0])
elif rotation == 8: # rotation 90
s = (s[1], s[0])
except:
pass
return s
def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix=''):
# Make sure only the first process in DDP process the dataset first, and the following others can use the cache
with torch_distributed_zero_first(rank):
dataset = LoadImagesAndLabels(path, imgsz, batch_size,
augment=augment, # augment images
hyp=hyp, # augmentation hyperparameters
rect=rect, # rectangular training
cache_images=cache,
single_cls=opt.single_cls,
stride=int(stride),
pad=pad,
image_weights=image_weights,
prefix=prefix)
batch_size = min(batch_size, len(dataset))
nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers]) # number of workers
sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None
loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader
# Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()
dataloader = loader(dataset,
batch_size=batch_size,
num_workers=0,
sampler=sampler,
pin_memory=True,
collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)
return dataloader, dataset
class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):
""" Dataloader that reuses workers
Uses same syntax as vanilla DataLoader
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
self.iterator = super().__iter__()
def __len__(self):
return len(self.batch_sampler.sampler)
def __iter__(self):
for i in range(len(self)):
yield next(self.iterator)
class _RepeatSampler(object):
""" Sampler that repeats forever
Args:
sampler (Sampler)
"""
def __init__(self, sampler):
self.sampler = sampler
def __iter__(self):
while True:
yield from iter(self.sampler)
class LoadImages: # for inference
def __init__(self, path, img_size=640, stride=32):
# p = str(Path(path).absolute()) # os-agnostic absolute path
p = str(Path(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.stride = stride
self.files = images + videos
self.nf = ni + nv # number of files
self.video_flag = [False] * ni + [True] * nv
self.mode = 'image'
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()
if not ret_val:
self.count += 1
self.cap.release()
if self.count == self.nf: # last video
raise StopIteration
else:
path = self.files[self.count]
self.new_video(path)
ret_val, img0 = self.cap.read()
self.frame += 1
print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='')
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='')
# Padded resize
img = letterbox(img0, self.img_size, stride=self.stride)[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
return path, img, img0, self.cap
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 LoadWebcam: # for inference
def __init__(self, pipe='0', img_size=640, stride=32):
self.img_size = img_size
self.stride = stride
if pipe.isnumeric():
pipe = eval(pipe) # local camera
# pipe = 'rtsp://192.168.1.64/1' # IP camera
# pipe = 'rtsp://username:[email protected]/1' # IP camera with login
# pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg' # IP golf camera
self.pipe = pipe
self.cap = cv2.VideoCapture(pipe) # video capture object
self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
def __iter__(self):
self.count = -1
return self
def __next__(self):
self.count += 1
if cv2.waitKey(1) == ord('q'): # q to quit
self.cap.release()
cv2.destroyAllWindows()
raise StopIteration
# Read frame
if self.pipe == 0: # local camera
ret_val, img0 = self.cap.read
基于深度学习的电动车头盔检测，使用yolo算法，并添加CBAM注意力机制提高检测的准确率。.zip
需积分: 5 12 浏览量
2024-05-10
09:40:51
上传
评论
收藏 2.22MB ZIP 举报
基于深度学习的电动车头盔检测，使用yolo算法，并添加CBAM注意力机制提高检测的准确率。.zip （2000个子文件） 
README.md 244B datasets.py 44KB general.py 28KB plots.py 18KB wandb_utils.py 16KB torch_utils.py 12KB loss.py 9KB metrics.py 9KB autoanchor.py 7KB hubconf.py 5KB google_utils.py 5KB count_size.py 4KB main_window.py 3KB main_window_ui.py 2KB activations.py 2KB resume.py 1KB log_dataset.py 819B class_num.py 718B __init__.py 0B __init__.py 0B __init__.py 0B userdata.sh 1KB mime.sh 780B download_weights.sh 277B
yolo-1615.txt 2KB yolo-623.txt 2KB yolo_107.txt 2KB yolo-1323.txt 1KB yolo-1253.txt 1KB yolo-496.txt 1KB yolo-369.txt 1KB yolo-937.txt 1KB yolo_41.txt 1KB yolo-429.txt 1KB yolo-446.txt 1KB yolo_16.txt 1KB yolo_108.txt 1KB yolo-427.txt 1KB yolo-590.txt 1KB yolo_37.txt 1KB yolo_35.txt 1KB yolo-457.txt 1KB yolo_80.txt 1KB yolo_52.txt 1KB yolo_106.txt 1KB yolo-498.txt 1KB yolo-959.txt 1KB yolo-381.txt 1KB yolo-461.txt 1020B yolo-230.txt 1014B yolo-269.txt 981B yolo-379.txt 967B yolo-271.txt 966B yolo-1595.txt 964B yolo-456.txt 962B yolo_150.txt 959B yolo-802.txt 958B yolo_102.txt 940B yolo_36.txt 936B yolo-613.txt 932B yolo-441.txt 928B yolo-630.txt 919B yolo-376.txt 916B yolo-202.txt 916B yolo-183.txt 912B yolo-462.txt 899B yolo-347.txt 886B yolo-313.txt 873B yolo-177.txt 862B yolo-1106.txt 860B yolo-422.txt 853B yolo-284.txt 848B yolo-957.txt 840B yolo_67.txt 832B yolo-1048.txt 826B yolo-563.txt 826B yolo-224.txt 825B yolo-1507.txt 823B yolo_113.txt 808B yolo_87.txt 806B yolo-312.txt 806B yolo-276.txt 803B yolo-1240.txt 802B yolo-210.txt 800B yolo-350.txt 797B yolo-259.txt 793B yolo-374.txt 790B yolo-436.txt 782B yolo-231.txt 772B yolo-870.txt 770B yolo-223.txt 760B yolo_105.txt 758B yolo_48.txt 754B yolo-561.txt 754B yolo_50.txt 749B yolo-199.txt 746B yolo-1566.txt 742B yolo-661.txt 735B yolo-1306.txt 729B yolo-1666.txt 720B共 2000 条
- 1
- 2
- 3
- 4
- 5
- 6
- 20
资源评论
生瓜蛋子
- 粉丝: 3809
- 资源: 4660
