From 7b6cba86efd6ee64bbe653c30276ce9c4608ebc1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 25 Jul 2019 13:19:26 +0200
Subject: [PATCH] updates

---
 models.py         |  4 ----
 train.py          |  4 ++--
 utils/adabound.py |  1 +
 utils/datasets.py |  5 ++---
 utils/utils.py    | 13 +++++--------
 5 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/models.py b/models.py
index 4e5d487c..a951a1a3 100755
--- a/models.py
+++ b/models.py
@@ -1,7 +1,3 @@
-import os
-
-import torch.nn.functional as F
-
 from utils.parse_config import *
 from utils.utils import *
 
diff --git a/train.py b/train.py
index 29afa67d..a734a6d5 100644
--- a/train.py
+++ b/train.py
@@ -1,15 +1,15 @@
 import argparse
 import time
 
+import torch.distributed as dist
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler
-import torch.distributed as dist
 
 import test  # import test.py to get mAP after each epoch
 from models import *
+from utils.adabound import *
 from utils.datasets import *
 from utils.utils import *
-from utils.adabound import *
 
 mixed_precision = True
 try:  # Mixed precision training https://github.com/NVIDIA/apex
diff --git a/utils/adabound.py b/utils/adabound.py
index 118e2a17..142b1262 100644
--- a/utils/adabound.py
+++ b/utils/adabound.py
@@ -1,4 +1,5 @@
 import math
+
 import torch
 from torch.optim import Optimizer
 
diff --git a/utils/datasets.py b/utils/datasets.py
index abf184a9..b4161b6b 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -8,9 +8,9 @@ from pathlib import Path
 import cv2
 import numpy as np
 import torch
+from PIL import Image, ExifTags
 from torch.utils.data import Dataset
 from tqdm import tqdm
-from PIL import Image, ExifTags
 
 from utils.utils import xyxy2xywh, xywh2xyxy
 
@@ -154,8 +154,7 @@ class LoadWebcam:  # for inference
 class LoadImagesAndLabels(Dataset):  # for training/testing
     def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False):
         with open(path, 'r') as f:
-            img_files = f.read().splitlines()
-            self.img_files = [x for x in img_files if os.path.splitext(x)[-1].lower() in img_formats]
+            self.img_files = [x for x in f.read().splitlines() if os.path.splitext(x)[-1].lower() in img_formats]
 
         n = len(self.img_files)
         bi = np.floor(np.arange(n) / batch_size).astype(np.int)  # batch index
diff --git a/utils/utils.py b/utils/utils.py
index c37f0af7..ec774d3a 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -1,5 +1,7 @@
 import glob
+import os
 import random
+from pathlib import Path
 
 import cv2
 import matplotlib
@@ -9,7 +11,6 @@ import torch
 import torch.nn as nn
 from PIL import Image
 from tqdm import tqdm
-from pathlib import Path
 
 from . import torch_utils  # , google_utils
 
@@ -543,18 +544,14 @@ def select_best_evolve(path='evolve*.txt'):  # from utils.utils import *; select
 
 
 def kmeans_targets(path='./data/coco_64img.txt'):  # from utils.utils import *; kmeans_targets()
+    img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif']
     with open(path, 'r') as f:
-        img_files = f.read().splitlines()
-        img_files = list(filter(lambda x: len(x) > 0, img_files))
+        img_files = [x for x in f.read().splitlines() if os.path.splitext(x)[-1].lower() in img_formats]
 
     # Read shapes
     n = len(img_files)
     assert n > 0, 'No images found in %s' % path
-    label_files = [x.replace('images', 'labels').
-                       replace('.jpeg', '.txt').
-                       replace('.jpg', '.txt').
-                       replace('.bmp', '.txt').
-                       replace('.png', '.txt') for x in img_files]
+    label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt') for x in img_files]
     s = np.array([Image.open(f).size for f in tqdm(img_files, desc='Reading image shapes')])  # (width, height)
 
     # Read targets