updates
This commit is contained in:
parent
e27b124828
commit
63c2736c12
20
test.py
20
test.py
|
@ -17,7 +17,8 @@ def test(cfg,
|
|||
conf_thres=0.001,
|
||||
nms_thres=0.5,
|
||||
save_json=False,
|
||||
model=None):
|
||||
model=None,
|
||||
dataloader=None):
|
||||
# Initialize/load model and set device
|
||||
if model is None:
|
||||
device = torch_utils.select_device(opt.device, batch_size=batch_size)
|
||||
|
@ -46,13 +47,14 @@ def test(cfg,
|
|||
names = load_classes(data['names']) # class names
|
||||
|
||||
# Dataloader
|
||||
dataset = LoadImagesAndLabels(test_path, img_size, batch_size)
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
dataloader = DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
num_workers=min([os.cpu_count(), batch_size if batch_size > 1 else 0, 16]),
|
||||
pin_memory=True,
|
||||
collate_fn=dataset.collate_fn)
|
||||
if dataloader is None:
|
||||
dataset = LoadImagesAndLabels(test_path, img_size, batch_size, rect=True)
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
dataloader = DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
num_workers=min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]),
|
||||
pin_memory=True,
|
||||
collate_fn=dataset.collate_fn)
|
||||
|
||||
seen = 0
|
||||
model.eval()
|
||||
|
@ -167,7 +169,7 @@ def test(cfg,
|
|||
|
||||
# Save JSON
|
||||
if save_json and map and len(jdict):
|
||||
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataset.img_files]
|
||||
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
|
||||
with open('results.json', 'w') as file:
|
||||
json.dump(jdict, file)
|
||||
|
||||
|
|
26
train.py
26
train.py
|
@ -72,6 +72,7 @@ def train():
|
|||
# Configure run
|
||||
data_dict = parse_data_cfg(data)
|
||||
train_path = data_dict['train']
|
||||
test_path = data_dict['valid']
|
||||
nc = int(data_dict['classes']) # number of classes
|
||||
|
||||
# Remove previous results
|
||||
|
@ -187,19 +188,17 @@ def train():
|
|||
model.yolo_layers = model.module.yolo_layers # move yolo layer indices to top level
|
||||
|
||||
# Dataset
|
||||
dataset = LoadImagesAndLabels(train_path,
|
||||
img_size,
|
||||
batch_size,
|
||||
dataset = LoadImagesAndLabels(train_path, img_size, batch_size,
|
||||
augment=True,
|
||||
hyp=hyp, # augmentation hyperparameters
|
||||
rect=opt.rect, # rectangular training
|
||||
image_weights=opt.img_weights,
|
||||
cache_labels=True if epochs > 10 else False,
|
||||
cache_images=False if opt.prebias else opt.cache_images)
|
||||
cache_labels=epochs > 10,
|
||||
cache_images=opt.cache_images and not opt.prebias)
|
||||
|
||||
# Dataloader
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 16]) # number of workers
|
||||
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
|
||||
print('Using %g dataloader workers' % nw)
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
|
@ -208,13 +207,23 @@ def train():
|
|||
pin_memory=True,
|
||||
collate_fn=dataset.collate_fn)
|
||||
|
||||
# Test Dataloader
|
||||
if not opt.prebias:
|
||||
testloader = torch.utils.data.DataLoader(LoadImagesAndLabels(test_path, img_size, batch_size, hyp=hyp,
|
||||
cache_labels=True,
|
||||
cache_images=opt.cache_images),
|
||||
batch_size=batch_size,
|
||||
num_workers=nw,
|
||||
pin_memory=True,
|
||||
collate_fn=dataset.collate_fn)
|
||||
|
||||
# Start training
|
||||
nb = len(dataloader)
|
||||
model.nc = nc # attach number of classes to model
|
||||
model.arc = opt.arc # attach yolo architecture
|
||||
model.hyp = hyp # attach hyperparameters to model
|
||||
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
|
||||
torch_utils.model_info(model, report='summary') # 'full' or 'summary'
|
||||
nb = len(dataloader)
|
||||
maps = np.zeros(nc) # mAP per class
|
||||
# torch.autograd.set_detect_anomaly(True)
|
||||
results = (0, 0, 0, 0, 0, 0, 0) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
|
||||
|
@ -321,7 +330,8 @@ def train():
|
|||
img_size=opt.img_size,
|
||||
model=model,
|
||||
conf_thres=0.001 if final_epoch and epoch > 0 else 0.1, # 0.1 for speed
|
||||
save_json=final_epoch and epoch > 0 and 'coco.data' in data)
|
||||
save_json=final_epoch and epoch > 0 and 'coco.data' in data,
|
||||
dataloader=testloader)
|
||||
|
||||
# Write epoch results
|
||||
with open(results_file, 'a') as f:
|
||||
|
|
|
@ -255,7 +255,7 @@ class LoadStreams: # multiple IP or RTSP cameras
|
|||
|
||||
|
||||
class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=True, image_weights=False,
|
||||
def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
|
||||
cache_labels=False, cache_images=False):
|
||||
path = str(Path(path)) # os-agnostic
|
||||
with open(path, 'r') as f:
|
||||
|
@ -319,7 +319,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
|||
self.labels = [np.zeros((0, 5))] * n
|
||||
extract_bounding_boxes = False
|
||||
create_datasubset = False
|
||||
pbar = tqdm(self.label_files, desc='Reading labels')
|
||||
pbar = tqdm(self.label_files, desc='Caching labels')
|
||||
nm, nf, ne, ns = 0, 0, 0, 0 # number missing, number found, number empty, number datasubset
|
||||
for i, file in enumerate(pbar):
|
||||
try:
|
||||
|
@ -370,13 +370,17 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
|||
ne += 1 # print('empty labels for image %s' % self.img_files[i]) # file empty
|
||||
# os.system("rm '%s' '%s'" % (self.img_files[i], self.label_files[i])) # remove
|
||||
|
||||
pbar.desc = 'Reading labels (%g found, %g missing, %g empty for %g images)' % (nf, nm, ne, n)
|
||||
pbar.desc = 'Caching labels (%g found, %g missing, %g empty for %g images)' % (nf, nm, ne, n)
|
||||
assert nf > 0, 'No labels found. Recommend correcting image and label paths.'
|
||||
|
||||
# Cache images into memory for faster training (~5GB)
|
||||
if cache_images and augment: # if training
|
||||
for i in tqdm(range(min(len(self.img_files), 10000)), desc='Reading images'): # max 10k images
|
||||
# Cache images into memory for faster training (WARNING: Large datasets may exceed system RAM)
|
||||
if cache_images: # if training
|
||||
gb = 0 # Gigabytes of cached images
|
||||
pbar = tqdm(range(len(self.img_files)), desc='Caching images')
|
||||
for i in pbar: # max 10k images
|
||||
self.imgs[i] = load_image(self, i)
|
||||
gb += self.imgs[i].nbytes
|
||||
pbar.desc = 'Caching images (%.1fGB)' % (gb / 1E9)
|
||||
|
||||
# Detect corrupted images https://medium.com/joelthchao/programmatically-detect-corrupted-image-8c1b2006c3d3
|
||||
detect_corrupted_images = False
|
||||
|
@ -503,10 +507,10 @@ def load_image(self, index):
|
|||
img_path = self.img_files[index]
|
||||
img = cv2.imread(img_path) # BGR
|
||||
assert img is not None, 'Image Not Found ' + img_path
|
||||
r = self.img_size / max(img.shape) # size ratio
|
||||
if self.augment: # if training (NOT testing), downsize to inference shape
|
||||
r = self.img_size / max(img.shape) # resize image to img_size
|
||||
if (r < 1) or ((r > 1) and self.augment): # always resize down, only resize up if training with augmentation
|
||||
h, w = img.shape[:2]
|
||||
img = cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_LINEAR) # _LINEAR fastest
|
||||
return cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_LINEAR) # _LINEAR fastest
|
||||
return img
|
||||
|
||||
|
||||
|
@ -569,13 +573,11 @@ def load_mosaic(self, index):
|
|||
# Concat/clip labels
|
||||
if len(labels4):
|
||||
labels4 = np.concatenate(labels4, 0)
|
||||
np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use before random_affine
|
||||
# np.clip(labels4[:, 1:], s / 2, 1.5 * s, out=labels4[:, 1:])
|
||||
# labels4[:, 1:] -= s / 2
|
||||
|
||||
# img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)]
|
||||
# np.clip(labels4[:, 1:] - s / 2, 0, s, out=labels4[:, 1:]) # use with center crop
|
||||
np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use with random_affine
|
||||
|
||||
# Augment
|
||||
# img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)] # center crop (WARNING, requires box pruning)
|
||||
img4, labels4 = random_affine(img4, labels4,
|
||||
degrees=self.hyp['degrees'],
|
||||
translate=self.hyp['translate'],
|
||||
|
|
Loading…
Reference in New Issue