From 3530aa5740860bcd48f460001de383c89aa41b0c Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Fri, 15 Jan 2021 16:51:44 +0000
Subject: [PATCH 01/36] Add helper functions for wandb and artifacts

---
 log_dataset.py       |  86 ++++++++++++++++++++
 utils/wandb_utils.py | 187 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 273 insertions(+)
 create mode 100644 log_dataset.py
 create mode 100644 utils/wandb_utils.py

diff --git a/log_dataset.py b/log_dataset.py
new file mode 100644
index 000000000000..5c6c0cf5576a
--- /dev/null
+++ b/log_dataset.py
@@ -0,0 +1,86 @@
+import argparse
+import logging
+import os
+import random
+import time
+from pathlib import Path
+from threading import Thread
+from warnings import warn
+
+import torch
+import yaml
+import wandb
+from utils.wandb_utils import WandbLogger
+
+from utils.general import check_dataset
+from utils.torch_utils import torch_distributed_zero_first
+from utils.datasets import create_dataloader
+
+WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+
+def create_dataset_artifact(opt):
+    with open(opt.data) as f:
+        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
+    wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
+    # Hyperparameters
+    with open(opt.hyp) as f:
+        hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
+        if 'box' not in hyp:
+            warn('Compatibility: %s missing "box" which was renamed from "giou" in %s' %
+                 (opt.hyp, 'https://github.com/ultralytics/yolov5/pull/1120'))
+            hyp['box'] = hyp.pop('giou')
+
+    with torch_distributed_zero_first(-1):
+        check_dataset(data_dict)  # check
+    train_path = data_dict['train']
+    test_path = data_dict['val']
+    nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
+    imgsz, batch_size = opt.img_size, opt.batch_size
+    assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
+    trainloader = create_dataloader(train_path, imgsz, batch_size, 32, opt,
+                                            hyp=hyp, cache=opt.cache_images, rect=opt.rect, rank=-1,
+                                            world_size=1, workers=opt.workers)[0]
+    
+    testloader = create_dataloader(test_path, imgsz, batch_size, 32, opt,  # testloader
+                                       hyp=hyp, cache=opt.cache_images, rect=True,
+                                       rank=-1, world_size=1, workers=opt.workers, pad=0.5)[0]
+    names_to_ids = {k: v for k, v in enumerate(names)}
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
+    wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
+    #Update/Create new config file with links to artifact
+    data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
+    data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
+    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.','_wandb.')
+    data_dict.pop('download',None) #Don't download the original dataset. Use artifacts
+    with open(ouput_data_config, 'w') as fp:
+        yaml.dump(data_dict, fp)
+    print("New Config file => ", ouput_data_config)
+    
+    
+    
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
+    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
+    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
+    parser.add_argument('--rect', action='store_true', help='rectangular training')
+    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
+    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
+    parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
+    parser.add_argument('--img-size', nargs='+', type=int, default=640, help='[train, test] image sizes')
+    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
+    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
+    parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
+    opt = parser.parse_args()
+
+    create_dataset_artifact(opt)
+        
+
+
+
+
+
+    
+    
diff --git a/utils/wandb_utils.py b/utils/wandb_utils.py
new file mode 100644
index 000000000000..827cbde21e27
--- /dev/null
+++ b/utils/wandb_utils.py
@@ -0,0 +1,187 @@
+from datetime import datetime
+import shutil
+from pathlib import Path
+import os
+import stat
+import logging
+
+import tqdm
+import torch
+import json
+from utils.general import xywh2xyxy
+logger = logging.getLogger(__name__)
+
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+    print("Install Weights & Biases for experiment logging via 'pip install wandb' (recommended)")
+
+WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+
+def remove_prefix(from_string, prefix):
+    return from_string[len(prefix):]
+    
+class WandbLogger():
+    def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
+        self.wandb = wandb
+        if self.wandb:
+            self.wandb_run = wandb.init(config=opt, resume="allow",
+                                   project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                   name=name,
+                                   job_type=job_type,
+                                   id=run_id)
+        else:
+            self.wandb_run = None
+        if job_type == 'Training':
+            self.setup_training(opt, data_dict)
+            if opt.bbox_interval == -1:
+                opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
+            if opt.save_period == -1:
+                opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
+    
+    def setup_training(self, opt, data_dict):
+        self.log_dict = {}
+        self.train_artifact_path, self.trainset_artifact = self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
+        self.test_artifact_path, self.testset_artifact = self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
+        self.result_artifact, self.result_table, self.weights = None, None, None
+        if self.train_artifact_path is not None:
+            train_path = self.train_artifact_path + '/data/images/'
+            data_dict['train'] = train_path
+        if self.test_artifact_path is not None:
+            test_path = self.test_artifact_path + '/data/images/'
+            data_dict['val'] = test_path
+            self.result_artifact = wandb.Artifact("run_"+wandb.run.id+"_progress", "evaluation")
+            self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
+        if opt.resume_from_artifact:
+            modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
+            if modeldir:
+                self.weights = modeldir + "/best.pt"
+                opt.weights = self.weights
+        
+
+    def download_dataset_artifact(self,path, alias):
+        if path.startswith(WANDB_ARTIFACT_PREFIX):
+            dataset_artifact = wandb.use_artifact(remove_prefix(path,WANDB_ARTIFACT_PREFIX)+":"+alias)
+            if dataset_artifact is None:
+                logger.error('Error: W&B dataset artifact doesn\'t exist')
+                raise ValueError('Artifact doesn\'t exist')
+            datadir = dataset_artifact.download()
+            labels_zip = datadir+"/data/labels.zip"
+            shutil.unpack_archive(labels_zip, datadir+'/data/labels', 'zip')
+            print("Downloaded dataset to : ", datadir)
+            return datadir, dataset_artifact
+        return None, None
+    
+    def download_model_artifact(self,name):
+        model_artifact = wandb.use_artifact(name+":latest")
+        if model_artifact is None:
+            logger.error('Error: W&B model artifact doesn\'t exist')
+            raise ValueError('Artifact doesn\'t exist')
+        modeldir = model_artifact.download()
+        print("Downloaded model to : ", modeldir)
+        return modeldir, model_artifact
+    
+    def log_model(self, path, opt, epoch):
+        datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
+        model_artifact = wandb.Artifact('run_'+wandb.run.id+'_model', type='model', metadata={
+            'original_url': str(path),
+            'epoch': epoch+1,
+            'save period': opt.save_period,
+            'project': opt.project,
+            'datetime': datetime_suffix
+        })
+        model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
+        model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
+        wandb.log_artifact(model_artifact)
+
+        if epoch+1 == opt.epochs:
+            model_artifact = wandb.Artifact('final_model', type='model', metadata={
+            'run_id': wandb.run.id,
+            'datetime': datetime_suffix
+            })
+            model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
+            model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
+            wandb.log_artifact(model_artifact)
+        print("Saving model artifact on epoch ", epoch+1)
+    
+    def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
+        artifact = wandb.Artifact(name=name, type="dataset")
+        image_path = dataloader.dataset.path
+        artifact.add_dir(image_path,name='data/images')
+        table = wandb.Table(
+            columns=["id", "train_image", "Classes"]
+        )
+        id_count = 0
+        class_set = wandb.Classes([{'id':id , 'name':name} for id,name in class_to_id.items()])
+        for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
+            targets = targets.to(device)
+            nb, _, height, width = img.shape  # batch size, channels, height, width
+            targets[:,2:] = (xywh2xyxy(targets[:,2:].view(-1, 4)))
+            for si, _ in enumerate(img):
+                height, width = shapes[si][0]
+                labels = targets[targets[:, 0] == si]
+                labels[:,2:] *=  torch.Tensor([width, height, width, height]).to(device)
+                labels = labels[:, 1:]
+                box_data = []
+                img_classes = {}
+                for cls, *xyxy in labels.tolist():
+                  class_id = int(cls)
+                  box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                  "class_id": class_id,
+                                  "box_caption": "%s" % (class_to_id[class_id]),
+                                  "scores": {"acc": 1},
+                                  "domain": "pixel"})
+                  img_classes[class_id] = class_to_id[class_id]
+                boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
+                table.add_data(id_count,wandb.Image(paths[si], classes=class_set, boxes=boxes), json.dumps(img_classes))
+                id_count = id_count+1
+        artifact.add(table, name)
+        label_path = image_path.replace('images','labels')
+        # Workaround for: Unable to log empty txt files via artifacts
+        if not os.path.isfile(name+'_labels.zip'): # make_archive won't check if file exists
+            shutil.make_archive(name+'_labels', 'zip', label_path)
+        artifact.add_file(name+'_labels.zip', name='data/labels.zip')
+        wandb.log_artifact(artifact)
+        print("Saving data to W&B...")
+            
+    def log(self, log_dict):
+        if self.wandb_run:
+            for key, value in log_dict.items():
+                self.log_dict[key] = value
+                
+    def end_epoch(self):
+        if self.wandb_run and self.log_dict:
+            wandb.log(self.log_dict)
+        self.log_dict = {}
+
+    def finish_run(self):
+        if self.wandb_run:
+            if self.result_artifact:
+                print("Add Training Progress Artifact")
+                self.result_artifact.add(self.result_table, 'result')
+                train_results = wandb.JoinedTable(self.testset_artifact.get("val"), self.result_table, "id")
+                self.result_artifact.add(train_results, 'joined_result')
+                wandb.log_artifact(self.result_artifact)
+            if self.log_dict:
+                wandb.log(self.log_dict)
+            wandb.run.finish()
+            
+    # !!!!! WIP !!!!
+    def add_to_training_progress(self, pred, class_map, class_dict_list, epoch, index):
+        # Painfully slow!!! investigate. 1) pred too large? replace class map dicts with calss variable?
+        if self.wandb_run and self.result_table and self.testset_artifact:
+            box_data = []
+            testset_table = self.testset_artifact.get("val").data
+            for *xyxy, conf, cls in pred.tolist():
+                if conf >= 0.175: # Arbitrary conf, make this an argparse
+                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                     "class_id": int(cls),
+                                     "box_caption": "%s %.3f" % (class_map[cls], conf),
+                                     "scores": {"class_score": conf},
+                                     "domain": "pixel"})
+            boxes = {"predictions": {"box_data": box_data, "class_labels": class_map}}  # inference-space
+            self.result_table.add_data(epoch,
+                                  index,
+                                  wandb.Image(testset_table[index][1], boxes=boxes, classes=class_dict_list))

From 1ef16b1edd99b87cbbb82f9675bcc95128976fcc Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Fri, 15 Jan 2021 17:05:36 +0000
Subject: [PATCH 02/36] cleanup

---
 utils/wandb_utils.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/utils/wandb_utils.py b/utils/wandb_utils.py
index 827cbde21e27..e2ebb2659138 100644
--- a/utils/wandb_utils.py
+++ b/utils/wandb_utils.py
@@ -167,21 +167,3 @@ def finish_run(self):
             if self.log_dict:
                 wandb.log(self.log_dict)
             wandb.run.finish()
-            
-    # !!!!! WIP !!!!
-    def add_to_training_progress(self, pred, class_map, class_dict_list, epoch, index):
-        # Painfully slow!!! investigate. 1) pred too large? replace class map dicts with calss variable?
-        if self.wandb_run and self.result_table and self.testset_artifact:
-            box_data = []
-            testset_table = self.testset_artifact.get("val").data
-            for *xyxy, conf, cls in pred.tolist():
-                if conf >= 0.175: # Arbitrary conf, make this an argparse
-                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                     "class_id": int(cls),
-                                     "box_caption": "%s %.3f" % (class_map[cls], conf),
-                                     "scores": {"class_score": conf},
-                                     "domain": "pixel"})
-            boxes = {"predictions": {"box_data": box_data, "class_labels": class_map}}  # inference-space
-            self.result_table.add_data(epoch,
-                                  index,
-                                  wandb.Image(testset_table[index][1], boxes=boxes, classes=class_dict_list))

From bd37ac778cf39359dde369a7a6f2c7b384e5b0b0 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Fri, 15 Jan 2021 21:51:40 +0000
Subject: [PATCH 03/36] Reorganize files

---
 log_dataset.py => utils/wandb/log_dataset.py | 0
 utils/{ => wandb}/wandb_utils.py             | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename log_dataset.py => utils/wandb/log_dataset.py (100%)
 rename utils/{ => wandb}/wandb_utils.py (100%)

diff --git a/log_dataset.py b/utils/wandb/log_dataset.py
similarity index 100%
rename from log_dataset.py
rename to utils/wandb/log_dataset.py
diff --git a/utils/wandb_utils.py b/utils/wandb/wandb_utils.py
similarity index 100%
rename from utils/wandb_utils.py
rename to utils/wandb/wandb_utils.py

From 7cd562336fb6b0c05ad36d306ea6d9769085c64b Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 17 Jan 2021 15:40:30 -0800
Subject: [PATCH 04/36] Update wandb_utils.py

---
 utils/wandb/wandb_utils.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/utils/wandb/wandb_utils.py b/utils/wandb/wandb_utils.py
index e2ebb2659138..774e01ecc059 100644
--- a/utils/wandb/wandb_utils.py
+++ b/utils/wandb/wandb_utils.py
@@ -8,15 +8,14 @@
 import tqdm
 import torch
 import json
-from utils.general import xywh2xyxy
+from utils.general import colorstr, xywh2xyxy
 logger = logging.getLogger(__name__)
 
-
 try:
     import wandb
 except ImportError:
     wandb = None
-    print("Install Weights & Biases for experiment logging via 'pip install wandb' (recommended)")
+    print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 

From d8daa648f357bbdc55240a9d0328f466ceefafc0 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 17 Jan 2021 15:42:05 -0800
Subject: [PATCH 05/36] Update log_dataset.py

We can remove this code, as the giou hyp has been deprecated for a while now.
---
 utils/wandb/log_dataset.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/utils/wandb/log_dataset.py b/utils/wandb/log_dataset.py
index 5c6c0cf5576a..f0866774ea98 100644
--- a/utils/wandb/log_dataset.py
+++ b/utils/wandb/log_dataset.py
@@ -25,10 +25,6 @@ def create_dataset_artifact(opt):
     # Hyperparameters
     with open(opt.hyp) as f:
         hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
-        if 'box' not in hyp:
-            warn('Compatibility: %s missing "box" which was renamed from "giou" in %s' %
-                 (opt.hyp, 'https://github.com/ultralytics/yolov5/pull/1120'))
-            hyp['box'] = hyp.pop('giou')
 
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check

From fc7d8f91c7edac3fdc20e6c4d2ed1c674ac73e59 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Mon, 18 Jan 2021 15:35:42 +0000
Subject: [PATCH 06/36] Reorganize and update dataloader call

---
 utils/datasets.py                             |  3 +-
 utils/wandb_logging/__init__.py               |  0
 utils/{wandb => wandb_logging}/log_dataset.py | 32 +++++++++----------
 utils/{wandb => wandb_logging}/wandb_utils.py |  4 +++
 4 files changed, 21 insertions(+), 18 deletions(-)
 create mode 100644 utils/wandb_logging/__init__.py
 rename utils/{wandb => wandb_logging}/log_dataset.py (73%)
 rename utils/{wandb => wandb_logging}/wandb_utils.py (98%)

diff --git a/utils/datasets.py b/utils/datasets.py
index 360d24c18874..1e23934b63cc 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -348,7 +348,8 @@ def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, r
         self.mosaic = self.augment and not self.rect  # load 4 images at a time into a mosaic (only during training)
         self.mosaic_border = [-img_size // 2, -img_size // 2]
         self.stride = stride
-
+        self.path = path
+        
         try:
             f = []  # image files
             for p in path if isinstance(path, list) else [path]:
diff --git a/utils/wandb_logging/__init__.py b/utils/wandb_logging/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/utils/wandb/log_dataset.py b/utils/wandb_logging/log_dataset.py
similarity index 73%
rename from utils/wandb/log_dataset.py
rename to utils/wandb_logging/log_dataset.py
index f0866774ea98..bacc025e0d55 100644
--- a/utils/wandb/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -7,12 +7,13 @@
 from threading import Thread
 from warnings import warn
 
+import numpy as np
 import torch
 import yaml
 import wandb
-from utils.wandb_utils import WandbLogger
+from wandb_utils import WandbLogger
 
-from utils.general import check_dataset
+from utils.general import check_dataset, colorstr
 from utils.torch_utils import torch_distributed_zero_first
 from utils.datasets import create_dataloader
 
@@ -33,13 +34,17 @@ def create_dataset_artifact(opt):
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
     imgsz, batch_size = opt.img_size, opt.batch_size
     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
-    trainloader = create_dataloader(train_path, imgsz, batch_size, 32, opt,
-                                            hyp=hyp, cache=opt.cache_images, rect=opt.rect, rank=-1,
-                                            world_size=1, workers=opt.workers)[0]
+    trainloader, trainset = create_dataloader(train_path, imgsz, batch_size, stride=32, opt=opt,
+                                            hyp=hyp, cache=False, rect=opt.rect, rank=-1,
+                                            world_size=opt.world_size, workers=opt.workers,
+                                            image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
+    mlc = np.concatenate(trainset.labels, 0)[:, 0].max()  # max label class
+    assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
     
-    testloader = create_dataloader(test_path, imgsz, batch_size, 32, opt,  # testloader
-                                       hyp=hyp, cache=opt.cache_images, rect=True,
-                                       rank=-1, world_size=1, workers=opt.workers, pad=0.5)[0]
+    testloader = create_dataloader(test_path, imgsz, batch_size, stride=32, opt=opt,  # testloader
+                                       hyp=hyp, cache=False, rect=True, rank=-1,
+                                       world_size=opt.world_size, workers=opt.workers,
+                                       pad=0.5, prefix=colorstr('val: '))[0]
     names_to_ids = {k: v for k, v in enumerate(names)}
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -62,21 +67,14 @@ def create_dataset_artifact(opt):
     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
     parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
     parser.add_argument('--rect', action='store_true', help='rectangular training')
-    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
     parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
     parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
     parser.add_argument('--img-size', nargs='+', type=int, default=640, help='[train, test] image sizes')
     parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
     parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
+    parser.add_argument('--quad', action='store_true', help='quad dataloader')
     parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
     opt = parser.parse_args()
+    opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
 
     create_dataset_artifact(opt)
-        
-
-
-
-
-
-    
-    
diff --git a/utils/wandb/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
similarity index 98%
rename from utils/wandb/wandb_utils.py
rename to utils/wandb_logging/wandb_utils.py
index 774e01ecc059..85b3439df299 100644
--- a/utils/wandb/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -2,12 +2,16 @@
 import shutil
 from pathlib import Path
 import os
+from os.path import dirname
 import stat
 import logging
 
 import tqdm
 import torch
 import json
+import sys
+# Add utils/ to path 
+sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))
 from utils.general import colorstr, xywh2xyxy
 logger = logging.getLogger(__name__)
 

From 9b8f46a3024ebab162989311dedda53072f5e8cc Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 20 Jan 2021 13:14:17 -0800
Subject: [PATCH 07/36] yaml.SafeLoader

---
 utils/wandb_logging/log_dataset.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index bacc025e0d55..a8e3eb9a7b54 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -21,11 +21,11 @@
 
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
-        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
+        data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
     wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
     # Hyperparameters
     with open(opt.hyp) as f:
-        hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
+        hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
 
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check

From 2c67ba11ca6cf5d57aa450accfbc3a12de880600 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 20 Jan 2021 13:22:25 -0800
Subject: [PATCH 08/36] PEP8 reformat

---
 utils/wandb_logging/log_dataset.py |  38 +++++-----
 utils/wandb_logging/wandb_utils.py | 110 +++++++++++++++--------------
 2 files changed, 73 insertions(+), 75 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index a8e3eb9a7b54..f609fbe6d80f 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,24 +1,18 @@
 import argparse
-import logging
 import os
-import random
-import time
-from pathlib import Path
-from threading import Thread
-from warnings import warn
 
 import numpy as np
 import torch
 import yaml
-import wandb
-from wandb_utils import WandbLogger
 
+from utils.datasets import create_dataloader
 from utils.general import check_dataset, colorstr
 from utils.torch_utils import torch_distributed_zero_first
-from utils.datasets import create_dataloader
+from wandb_utils import WandbLogger
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
+
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
@@ -35,32 +29,32 @@ def create_dataset_artifact(opt):
     imgsz, batch_size = opt.img_size, opt.batch_size
     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
     trainloader, trainset = create_dataloader(train_path, imgsz, batch_size, stride=32, opt=opt,
-                                            hyp=hyp, cache=False, rect=opt.rect, rank=-1,
-                                            world_size=opt.world_size, workers=opt.workers,
-                                            image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
+                                              hyp=hyp, cache=False, rect=opt.rect, rank=-1,
+                                              world_size=opt.world_size, workers=opt.workers,
+                                              image_weights=opt.image_weights, quad=opt.quad,
+                                              prefix=colorstr('train: '))
     mlc = np.concatenate(trainset.labels, 0)[:, 0].max()  # max label class
     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
-    
+
     testloader = create_dataloader(test_path, imgsz, batch_size, stride=32, opt=opt,  # testloader
-                                       hyp=hyp, cache=False, rect=True, rank=-1,
-                                       world_size=opt.world_size, workers=opt.workers,
-                                       pad=0.5, prefix=colorstr('val: '))[0]
+                                   hyp=hyp, cache=False, rect=True, rank=-1,
+                                   world_size=opt.world_size, workers=opt.workers,
+                                   pad=0.5, prefix=colorstr('val: '))[0]
     names_to_ids = {k: v for k, v in enumerate(names)}
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
     wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
-    #Update/Create new config file with links to artifact
+    # Update/Create new config file with links to artifact
     data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
     data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
-    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.','_wandb.')
-    data_dict.pop('download',None) #Don't download the original dataset. Use artifacts
+    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')
+    data_dict.pop('download', None)  # Don't download the original dataset. Use artifacts
     with open(ouput_data_config, 'w') as fp:
         yaml.dump(data_dict, fp)
     print("New Config file => ", ouput_data_config)
-    
-    
-    
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 85b3439df299..eac23ef9e841 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -1,18 +1,18 @@
-from datetime import datetime
-import shutil
-from pathlib import Path
+import json
+import logging
 import os
+import shutil
+import sys
+from datetime import datetime
 from os.path import dirname
-import stat
-import logging
+from pathlib import Path
 
-import tqdm
 import torch
-import json
-import sys
-# Add utils/ to path 
+
+# Add utils/ to path
 sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))
 from utils.general import colorstr, xywh2xyxy
+
 logger = logging.getLogger(__name__)
 
 try:
@@ -23,18 +23,20 @@
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
+
 def remove_prefix(from_string, prefix):
     return from_string[len(prefix):]
-    
+
+
 class WandbLogger():
     def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
         self.wandb = wandb
         if self.wandb:
             self.wandb_run = wandb.init(config=opt, resume="allow",
-                                   project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
-                                   name=name,
-                                   job_type=job_type,
-                                   id=run_id)
+                                        project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                        name=name,
+                                        job_type=job_type,
+                                        id=run_id)
         else:
             self.wandb_run = None
         if job_type == 'Training':
@@ -43,11 +45,13 @@ def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
                 opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
             if opt.save_period == -1:
                 opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
-    
+
     def setup_training(self, opt, data_dict):
         self.log_dict = {}
-        self.train_artifact_path, self.trainset_artifact = self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
-        self.test_artifact_path, self.testset_artifact = self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
+        self.train_artifact_path, self.trainset_artifact = \
+            self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
+        self.test_artifact_path, self.testset_artifact = \
+            self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
         self.result_artifact, self.result_table, self.weights = None, None, None
         if self.train_artifact_path is not None:
             train_path = self.train_artifact_path + '/data/images/'
@@ -55,42 +59,41 @@ def setup_training(self, opt, data_dict):
         if self.test_artifact_path is not None:
             test_path = self.test_artifact_path + '/data/images/'
             data_dict['val'] = test_path
-            self.result_artifact = wandb.Artifact("run_"+wandb.run.id+"_progress", "evaluation")
+            self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
             self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
         if opt.resume_from_artifact:
             modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
             if modeldir:
                 self.weights = modeldir + "/best.pt"
                 opt.weights = self.weights
-        
 
-    def download_dataset_artifact(self,path, alias):
+    def download_dataset_artifact(self, path, alias):
         if path.startswith(WANDB_ARTIFACT_PREFIX):
-            dataset_artifact = wandb.use_artifact(remove_prefix(path,WANDB_ARTIFACT_PREFIX)+":"+alias)
+            dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
             if dataset_artifact is None:
                 logger.error('Error: W&B dataset artifact doesn\'t exist')
                 raise ValueError('Artifact doesn\'t exist')
             datadir = dataset_artifact.download()
-            labels_zip = datadir+"/data/labels.zip"
-            shutil.unpack_archive(labels_zip, datadir+'/data/labels', 'zip')
+            labels_zip = datadir + "/data/labels.zip"
+            shutil.unpack_archive(labels_zip, datadir + '/data/labels', 'zip')
             print("Downloaded dataset to : ", datadir)
             return datadir, dataset_artifact
         return None, None
-    
-    def download_model_artifact(self,name):
-        model_artifact = wandb.use_artifact(name+":latest")
+
+    def download_model_artifact(self, name):
+        model_artifact = wandb.use_artifact(name + ":latest")
         if model_artifact is None:
             logger.error('Error: W&B model artifact doesn\'t exist')
             raise ValueError('Artifact doesn\'t exist')
         modeldir = model_artifact.download()
         print("Downloaded model to : ", modeldir)
         return modeldir, model_artifact
-    
+
     def log_model(self, path, opt, epoch):
         datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
-        model_artifact = wandb.Artifact('run_'+wandb.run.id+'_model', type='model', metadata={
+        model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
             'original_url': str(path),
-            'epoch': epoch+1,
+            'epoch': epoch + 1,
             'save period': opt.save_period,
             'project': opt.project,
             'datetime': datetime_suffix
@@ -99,61 +102,62 @@ def log_model(self, path, opt, epoch):
         model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
         wandb.log_artifact(model_artifact)
 
-        if epoch+1 == opt.epochs:
+        if epoch + 1 == opt.epochs:
             model_artifact = wandb.Artifact('final_model', type='model', metadata={
-            'run_id': wandb.run.id,
-            'datetime': datetime_suffix
+                'run_id': wandb.run.id,
+                'datetime': datetime_suffix
             })
             model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
             model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
             wandb.log_artifact(model_artifact)
-        print("Saving model artifact on epoch ", epoch+1)
-    
+        print("Saving model artifact on epoch ", epoch + 1)
+
     def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
         image_path = dataloader.dataset.path
-        artifact.add_dir(image_path,name='data/images')
+        artifact.add_dir(image_path, name='data/images')
         table = wandb.Table(
             columns=["id", "train_image", "Classes"]
         )
         id_count = 0
-        class_set = wandb.Classes([{'id':id , 'name':name} for id,name in class_to_id.items()])
+        class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
         for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
             targets = targets.to(device)
             nb, _, height, width = img.shape  # batch size, channels, height, width
-            targets[:,2:] = (xywh2xyxy(targets[:,2:].view(-1, 4)))
+            targets[:, 2:] = (xywh2xyxy(targets[:, 2:].view(-1, 4)))
             for si, _ in enumerate(img):
                 height, width = shapes[si][0]
                 labels = targets[targets[:, 0] == si]
-                labels[:,2:] *=  torch.Tensor([width, height, width, height]).to(device)
+                labels[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)
                 labels = labels[:, 1:]
                 box_data = []
                 img_classes = {}
                 for cls, *xyxy in labels.tolist():
-                  class_id = int(cls)
-                  box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                  "class_id": class_id,
-                                  "box_caption": "%s" % (class_to_id[class_id]),
-                                  "scores": {"acc": 1},
-                                  "domain": "pixel"})
-                  img_classes[class_id] = class_to_id[class_id]
+                    class_id = int(cls)
+                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                     "class_id": class_id,
+                                     "box_caption": "%s" % (class_to_id[class_id]),
+                                     "scores": {"acc": 1},
+                                     "domain": "pixel"})
+                    img_classes[class_id] = class_to_id[class_id]
                 boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-                table.add_data(id_count,wandb.Image(paths[si], classes=class_set, boxes=boxes), json.dumps(img_classes))
-                id_count = id_count+1
+                table.add_data(id_count, wandb.Image(paths[si], classes=class_set, boxes=boxes),
+                               json.dumps(img_classes))
+                id_count = id_count + 1
         artifact.add(table, name)
-        label_path = image_path.replace('images','labels')
+        label_path = image_path.replace('images', 'labels')
         # Workaround for: Unable to log empty txt files via artifacts
-        if not os.path.isfile(name+'_labels.zip'): # make_archive won't check if file exists
-            shutil.make_archive(name+'_labels', 'zip', label_path)
-        artifact.add_file(name+'_labels.zip', name='data/labels.zip')
+        if not os.path.isfile(name + '_labels.zip'):  # make_archive won't check if file exists
+            shutil.make_archive(name + '_labels', 'zip', label_path)
+        artifact.add_file(name + '_labels.zip', name='data/labels.zip')
         wandb.log_artifact(artifact)
         print("Saving data to W&B...")
-            
+
     def log(self, log_dict):
         if self.wandb_run:
             for key, value in log_dict.items():
                 self.log_dict[key] = value
-                
+
     def end_epoch(self):
         if self.wandb_run and self.log_dict:
             wandb.log(self.log_dict)

From 9ffb88727a3c89a523023cbb9a9ff40a409d1b42 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 20 Jan 2021 13:28:29 -0800
Subject: [PATCH 09/36] remove redundant checks

---
 utils/wandb_logging/log_dataset.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index f609fbe6d80f..31633af635d0 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,7 +1,6 @@
 import argparse
 import os
 
-import numpy as np
 import torch
 import yaml
 
@@ -20,23 +19,18 @@ def create_dataset_artifact(opt):
     # Hyperparameters
     with open(opt.hyp) as f:
         hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
-
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check
-    train_path = data_dict['train']
-    test_path = data_dict['val']
+
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
     imgsz, batch_size = opt.img_size, opt.batch_size
-    assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
-    trainloader, trainset = create_dataloader(train_path, imgsz, batch_size, stride=32, opt=opt,
+    trainloader, trainset = create_dataloader(data_dict['train'], imgsz, batch_size, stride=32, opt=opt,
                                               hyp=hyp, cache=False, rect=opt.rect, rank=-1,
                                               world_size=opt.world_size, workers=opt.workers,
                                               image_weights=opt.image_weights, quad=opt.quad,
                                               prefix=colorstr('train: '))
-    mlc = np.concatenate(trainset.labels, 0)[:, 0].max()  # max label class
-    assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
 
-    testloader = create_dataloader(test_path, imgsz, batch_size, stride=32, opt=opt,  # testloader
+    testloader = create_dataloader(data_dict['val'], imgsz, batch_size, stride=32, opt=opt,  # testloader
                                    hyp=hyp, cache=False, rect=True, rank=-1,
                                    world_size=opt.world_size, workers=opt.workers,
                                    pad=0.5, prefix=colorstr('val: '))[0]

From 1eca722fdcfdb71de835505c27d1ebbe8c494b72 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Fri, 15 Jan 2021 16:51:44 +0000
Subject: [PATCH 10/36] Add helper functions for wandb and artifacts

---
 log_dataset.py       |  86 ++++++++++++++++++++
 utils/wandb_utils.py | 187 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 273 insertions(+)
 create mode 100644 log_dataset.py
 create mode 100644 utils/wandb_utils.py

diff --git a/log_dataset.py b/log_dataset.py
new file mode 100644
index 000000000000..5c6c0cf5576a
--- /dev/null
+++ b/log_dataset.py
@@ -0,0 +1,86 @@
+import argparse
+import logging
+import os
+import random
+import time
+from pathlib import Path
+from threading import Thread
+from warnings import warn
+
+import torch
+import yaml
+import wandb
+from utils.wandb_utils import WandbLogger
+
+from utils.general import check_dataset
+from utils.torch_utils import torch_distributed_zero_first
+from utils.datasets import create_dataloader
+
+WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+
+def create_dataset_artifact(opt):
+    with open(opt.data) as f:
+        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
+    wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
+    # Hyperparameters
+    with open(opt.hyp) as f:
+        hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
+        if 'box' not in hyp:
+            warn('Compatibility: %s missing "box" which was renamed from "giou" in %s' %
+                 (opt.hyp, 'https://github.com/ultralytics/yolov5/pull/1120'))
+            hyp['box'] = hyp.pop('giou')
+
+    with torch_distributed_zero_first(-1):
+        check_dataset(data_dict)  # check
+    train_path = data_dict['train']
+    test_path = data_dict['val']
+    nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
+    imgsz, batch_size = opt.img_size, opt.batch_size
+    assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
+    trainloader = create_dataloader(train_path, imgsz, batch_size, 32, opt,
+                                            hyp=hyp, cache=opt.cache_images, rect=opt.rect, rank=-1,
+                                            world_size=1, workers=opt.workers)[0]
+    
+    testloader = create_dataloader(test_path, imgsz, batch_size, 32, opt,  # testloader
+                                       hyp=hyp, cache=opt.cache_images, rect=True,
+                                       rank=-1, world_size=1, workers=opt.workers, pad=0.5)[0]
+    names_to_ids = {k: v for k, v in enumerate(names)}
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
+    wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
+    #Update/Create new config file with links to artifact
+    data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
+    data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
+    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.','_wandb.')
+    data_dict.pop('download',None) #Don't download the original dataset. Use artifacts
+    with open(ouput_data_config, 'w') as fp:
+        yaml.dump(data_dict, fp)
+    print("New Config file => ", ouput_data_config)
+    
+    
+    
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
+    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
+    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
+    parser.add_argument('--rect', action='store_true', help='rectangular training')
+    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
+    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
+    parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
+    parser.add_argument('--img-size', nargs='+', type=int, default=640, help='[train, test] image sizes')
+    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
+    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
+    parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
+    opt = parser.parse_args()
+
+    create_dataset_artifact(opt)
+        
+
+
+
+
+
+    
+    
diff --git a/utils/wandb_utils.py b/utils/wandb_utils.py
new file mode 100644
index 000000000000..827cbde21e27
--- /dev/null
+++ b/utils/wandb_utils.py
@@ -0,0 +1,187 @@
+from datetime import datetime
+import shutil
+from pathlib import Path
+import os
+import stat
+import logging
+
+import tqdm
+import torch
+import json
+from utils.general import xywh2xyxy
+logger = logging.getLogger(__name__)
+
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+    print("Install Weights & Biases for experiment logging via 'pip install wandb' (recommended)")
+
+WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+
+def remove_prefix(from_string, prefix):
+    return from_string[len(prefix):]
+    
+class WandbLogger():
+    def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
+        self.wandb = wandb
+        if self.wandb:
+            self.wandb_run = wandb.init(config=opt, resume="allow",
+                                   project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                   name=name,
+                                   job_type=job_type,
+                                   id=run_id)
+        else:
+            self.wandb_run = None
+        if job_type == 'Training':
+            self.setup_training(opt, data_dict)
+            if opt.bbox_interval == -1:
+                opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
+            if opt.save_period == -1:
+                opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
+    
+    def setup_training(self, opt, data_dict):
+        self.log_dict = {}
+        self.train_artifact_path, self.trainset_artifact = self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
+        self.test_artifact_path, self.testset_artifact = self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
+        self.result_artifact, self.result_table, self.weights = None, None, None
+        if self.train_artifact_path is not None:
+            train_path = self.train_artifact_path + '/data/images/'
+            data_dict['train'] = train_path
+        if self.test_artifact_path is not None:
+            test_path = self.test_artifact_path + '/data/images/'
+            data_dict['val'] = test_path
+            self.result_artifact = wandb.Artifact("run_"+wandb.run.id+"_progress", "evaluation")
+            self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
+        if opt.resume_from_artifact:
+            modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
+            if modeldir:
+                self.weights = modeldir + "/best.pt"
+                opt.weights = self.weights
+        
+
+    def download_dataset_artifact(self,path, alias):
+        if path.startswith(WANDB_ARTIFACT_PREFIX):
+            dataset_artifact = wandb.use_artifact(remove_prefix(path,WANDB_ARTIFACT_PREFIX)+":"+alias)
+            if dataset_artifact is None:
+                logger.error('Error: W&B dataset artifact doesn\'t exist')
+                raise ValueError('Artifact doesn\'t exist')
+            datadir = dataset_artifact.download()
+            labels_zip = datadir+"/data/labels.zip"
+            shutil.unpack_archive(labels_zip, datadir+'/data/labels', 'zip')
+            print("Downloaded dataset to : ", datadir)
+            return datadir, dataset_artifact
+        return None, None
+    
+    def download_model_artifact(self,name):
+        model_artifact = wandb.use_artifact(name+":latest")
+        if model_artifact is None:
+            logger.error('Error: W&B model artifact doesn\'t exist')
+            raise ValueError('Artifact doesn\'t exist')
+        modeldir = model_artifact.download()
+        print("Downloaded model to : ", modeldir)
+        return modeldir, model_artifact
+    
+    def log_model(self, path, opt, epoch):
+        datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
+        model_artifact = wandb.Artifact('run_'+wandb.run.id+'_model', type='model', metadata={
+            'original_url': str(path),
+            'epoch': epoch+1,
+            'save period': opt.save_period,
+            'project': opt.project,
+            'datetime': datetime_suffix
+        })
+        model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
+        model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
+        wandb.log_artifact(model_artifact)
+
+        if epoch+1 == opt.epochs:
+            model_artifact = wandb.Artifact('final_model', type='model', metadata={
+            'run_id': wandb.run.id,
+            'datetime': datetime_suffix
+            })
+            model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
+            model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
+            wandb.log_artifact(model_artifact)
+        print("Saving model artifact on epoch ", epoch+1)
+    
+    def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
+        artifact = wandb.Artifact(name=name, type="dataset")
+        image_path = dataloader.dataset.path
+        artifact.add_dir(image_path,name='data/images')
+        table = wandb.Table(
+            columns=["id", "train_image", "Classes"]
+        )
+        id_count = 0
+        class_set = wandb.Classes([{'id':id , 'name':name} for id,name in class_to_id.items()])
+        for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
+            targets = targets.to(device)
+            nb, _, height, width = img.shape  # batch size, channels, height, width
+            targets[:,2:] = (xywh2xyxy(targets[:,2:].view(-1, 4)))
+            for si, _ in enumerate(img):
+                height, width = shapes[si][0]
+                labels = targets[targets[:, 0] == si]
+                labels[:,2:] *=  torch.Tensor([width, height, width, height]).to(device)
+                labels = labels[:, 1:]
+                box_data = []
+                img_classes = {}
+                for cls, *xyxy in labels.tolist():
+                  class_id = int(cls)
+                  box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                  "class_id": class_id,
+                                  "box_caption": "%s" % (class_to_id[class_id]),
+                                  "scores": {"acc": 1},
+                                  "domain": "pixel"})
+                  img_classes[class_id] = class_to_id[class_id]
+                boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
+                table.add_data(id_count,wandb.Image(paths[si], classes=class_set, boxes=boxes), json.dumps(img_classes))
+                id_count = id_count+1
+        artifact.add(table, name)
+        label_path = image_path.replace('images','labels')
+        # Workaround for: Unable to log empty txt files via artifacts
+        if not os.path.isfile(name+'_labels.zip'): # make_archive won't check if file exists
+            shutil.make_archive(name+'_labels', 'zip', label_path)
+        artifact.add_file(name+'_labels.zip', name='data/labels.zip')
+        wandb.log_artifact(artifact)
+        print("Saving data to W&B...")
+            
+    def log(self, log_dict):
+        if self.wandb_run:
+            for key, value in log_dict.items():
+                self.log_dict[key] = value
+                
+    def end_epoch(self):
+        if self.wandb_run and self.log_dict:
+            wandb.log(self.log_dict)
+        self.log_dict = {}
+
+    def finish_run(self):
+        if self.wandb_run:
+            if self.result_artifact:
+                print("Add Training Progress Artifact")
+                self.result_artifact.add(self.result_table, 'result')
+                train_results = wandb.JoinedTable(self.testset_artifact.get("val"), self.result_table, "id")
+                self.result_artifact.add(train_results, 'joined_result')
+                wandb.log_artifact(self.result_artifact)
+            if self.log_dict:
+                wandb.log(self.log_dict)
+            wandb.run.finish()
+            
+    # !!!!! WIP !!!!
+    def add_to_training_progress(self, pred, class_map, class_dict_list, epoch, index):
+        # Painfully slow!!! investigate. 1) pred too large? replace class map dicts with calss variable?
+        if self.wandb_run and self.result_table and self.testset_artifact:
+            box_data = []
+            testset_table = self.testset_artifact.get("val").data
+            for *xyxy, conf, cls in pred.tolist():
+                if conf >= 0.175: # Arbitrary conf, make this an argparse
+                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                     "class_id": int(cls),
+                                     "box_caption": "%s %.3f" % (class_map[cls], conf),
+                                     "scores": {"class_score": conf},
+                                     "domain": "pixel"})
+            boxes = {"predictions": {"box_data": box_data, "class_labels": class_map}}  # inference-space
+            self.result_table.add_data(epoch,
+                                  index,
+                                  wandb.Image(testset_table[index][1], boxes=boxes, classes=class_dict_list))

From ad8408fb99002c820468a017e46aa97e6b685131 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Fri, 15 Jan 2021 17:05:36 +0000
Subject: [PATCH 11/36] cleanup

---
 utils/wandb_utils.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/utils/wandb_utils.py b/utils/wandb_utils.py
index 827cbde21e27..e2ebb2659138 100644
--- a/utils/wandb_utils.py
+++ b/utils/wandb_utils.py
@@ -167,21 +167,3 @@ def finish_run(self):
             if self.log_dict:
                 wandb.log(self.log_dict)
             wandb.run.finish()
-            
-    # !!!!! WIP !!!!
-    def add_to_training_progress(self, pred, class_map, class_dict_list, epoch, index):
-        # Painfully slow!!! investigate. 1) pred too large? replace class map dicts with calss variable?
-        if self.wandb_run and self.result_table and self.testset_artifact:
-            box_data = []
-            testset_table = self.testset_artifact.get("val").data
-            for *xyxy, conf, cls in pred.tolist():
-                if conf >= 0.175: # Arbitrary conf, make this an argparse
-                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                     "class_id": int(cls),
-                                     "box_caption": "%s %.3f" % (class_map[cls], conf),
-                                     "scores": {"class_score": conf},
-                                     "domain": "pixel"})
-            boxes = {"predictions": {"box_data": box_data, "class_labels": class_map}}  # inference-space
-            self.result_table.add_data(epoch,
-                                  index,
-                                  wandb.Image(testset_table[index][1], boxes=boxes, classes=class_dict_list))

From 8b039b3cde8de7327f94b9e8ad0fe93ba6c73afb Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Fri, 15 Jan 2021 21:51:40 +0000
Subject: [PATCH 12/36] Reorganize files

---
 log_dataset.py => utils/wandb/log_dataset.py | 0
 utils/{ => wandb}/wandb_utils.py             | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename log_dataset.py => utils/wandb/log_dataset.py (100%)
 rename utils/{ => wandb}/wandb_utils.py (100%)

diff --git a/log_dataset.py b/utils/wandb/log_dataset.py
similarity index 100%
rename from log_dataset.py
rename to utils/wandb/log_dataset.py
diff --git a/utils/wandb_utils.py b/utils/wandb/wandb_utils.py
similarity index 100%
rename from utils/wandb_utils.py
rename to utils/wandb/wandb_utils.py

From 0b715eb904c14f49719b6902beeca85ef73ebaa4 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 17 Jan 2021 15:40:30 -0800
Subject: [PATCH 13/36] Update wandb_utils.py

---
 utils/wandb/wandb_utils.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/utils/wandb/wandb_utils.py b/utils/wandb/wandb_utils.py
index e2ebb2659138..774e01ecc059 100644
--- a/utils/wandb/wandb_utils.py
+++ b/utils/wandb/wandb_utils.py
@@ -8,15 +8,14 @@
 import tqdm
 import torch
 import json
-from utils.general import xywh2xyxy
+from utils.general import colorstr, xywh2xyxy
 logger = logging.getLogger(__name__)
 
-
 try:
     import wandb
 except ImportError:
     wandb = None
-    print("Install Weights & Biases for experiment logging via 'pip install wandb' (recommended)")
+    print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 

From 58ae23fa4355c5a3cc326c32525d71b237352582 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 17 Jan 2021 15:42:05 -0800
Subject: [PATCH 14/36] Update log_dataset.py

We can remove this code, as the giou hyp has been deprecated for a while now.
---
 utils/wandb/log_dataset.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/utils/wandb/log_dataset.py b/utils/wandb/log_dataset.py
index 5c6c0cf5576a..f0866774ea98 100644
--- a/utils/wandb/log_dataset.py
+++ b/utils/wandb/log_dataset.py
@@ -25,10 +25,6 @@ def create_dataset_artifact(opt):
     # Hyperparameters
     with open(opt.hyp) as f:
         hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
-        if 'box' not in hyp:
-            warn('Compatibility: %s missing "box" which was renamed from "giou" in %s' %
-                 (opt.hyp, 'https://github.com/ultralytics/yolov5/pull/1120'))
-            hyp['box'] = hyp.pop('giou')
 
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check

From 6671347609f32fec7d7ac6da2ca60b734633e416 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Mon, 18 Jan 2021 15:35:42 +0000
Subject: [PATCH 15/36] Reorganize and update dataloader call

---
 utils/wandb/log_dataset.py         |  82 --------------
 utils/wandb/wandb_utils.py         | 168 -----------------------------
 utils/wandb_logging/log_dataset.py |  54 ++++++----
 utils/wandb_logging/wandb_utils.py | 110 +++++++++----------
 4 files changed, 86 insertions(+), 328 deletions(-)
 delete mode 100644 utils/wandb/log_dataset.py
 delete mode 100644 utils/wandb/wandb_utils.py

diff --git a/utils/wandb/log_dataset.py b/utils/wandb/log_dataset.py
deleted file mode 100644
index f0866774ea98..000000000000
--- a/utils/wandb/log_dataset.py
+++ /dev/null
@@ -1,82 +0,0 @@
-import argparse
-import logging
-import os
-import random
-import time
-from pathlib import Path
-from threading import Thread
-from warnings import warn
-
-import torch
-import yaml
-import wandb
-from utils.wandb_utils import WandbLogger
-
-from utils.general import check_dataset
-from utils.torch_utils import torch_distributed_zero_first
-from utils.datasets import create_dataloader
-
-WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
-
-def create_dataset_artifact(opt):
-    with open(opt.data) as f:
-        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
-    wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
-    # Hyperparameters
-    with open(opt.hyp) as f:
-        hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
-
-    with torch_distributed_zero_first(-1):
-        check_dataset(data_dict)  # check
-    train_path = data_dict['train']
-    test_path = data_dict['val']
-    nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
-    imgsz, batch_size = opt.img_size, opt.batch_size
-    assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
-    trainloader = create_dataloader(train_path, imgsz, batch_size, 32, opt,
-                                            hyp=hyp, cache=opt.cache_images, rect=opt.rect, rank=-1,
-                                            world_size=1, workers=opt.workers)[0]
-    
-    testloader = create_dataloader(test_path, imgsz, batch_size, 32, opt,  # testloader
-                                       hyp=hyp, cache=opt.cache_images, rect=True,
-                                       rank=-1, world_size=1, workers=opt.workers, pad=0.5)[0]
-    names_to_ids = {k: v for k, v in enumerate(names)}
-
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
-    wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
-    #Update/Create new config file with links to artifact
-    data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
-    data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
-    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.','_wandb.')
-    data_dict.pop('download',None) #Don't download the original dataset. Use artifacts
-    with open(ouput_data_config, 'w') as fp:
-        yaml.dump(data_dict, fp)
-    print("New Config file => ", ouput_data_config)
-    
-    
-    
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
-    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
-    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
-    parser.add_argument('--rect', action='store_true', help='rectangular training')
-    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
-    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
-    parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
-    parser.add_argument('--img-size', nargs='+', type=int, default=640, help='[train, test] image sizes')
-    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
-    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
-    parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
-    opt = parser.parse_args()
-
-    create_dataset_artifact(opt)
-        
-
-
-
-
-
-    
-    
diff --git a/utils/wandb/wandb_utils.py b/utils/wandb/wandb_utils.py
deleted file mode 100644
index 774e01ecc059..000000000000
--- a/utils/wandb/wandb_utils.py
+++ /dev/null
@@ -1,168 +0,0 @@
-from datetime import datetime
-import shutil
-from pathlib import Path
-import os
-import stat
-import logging
-
-import tqdm
-import torch
-import json
-from utils.general import colorstr, xywh2xyxy
-logger = logging.getLogger(__name__)
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-    print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
-
-WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
-
-def remove_prefix(from_string, prefix):
-    return from_string[len(prefix):]
-    
-class WandbLogger():
-    def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
-        self.wandb = wandb
-        if self.wandb:
-            self.wandb_run = wandb.init(config=opt, resume="allow",
-                                   project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
-                                   name=name,
-                                   job_type=job_type,
-                                   id=run_id)
-        else:
-            self.wandb_run = None
-        if job_type == 'Training':
-            self.setup_training(opt, data_dict)
-            if opt.bbox_interval == -1:
-                opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
-            if opt.save_period == -1:
-                opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
-    
-    def setup_training(self, opt, data_dict):
-        self.log_dict = {}
-        self.train_artifact_path, self.trainset_artifact = self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
-        self.test_artifact_path, self.testset_artifact = self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
-        self.result_artifact, self.result_table, self.weights = None, None, None
-        if self.train_artifact_path is not None:
-            train_path = self.train_artifact_path + '/data/images/'
-            data_dict['train'] = train_path
-        if self.test_artifact_path is not None:
-            test_path = self.test_artifact_path + '/data/images/'
-            data_dict['val'] = test_path
-            self.result_artifact = wandb.Artifact("run_"+wandb.run.id+"_progress", "evaluation")
-            self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
-        if opt.resume_from_artifact:
-            modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
-            if modeldir:
-                self.weights = modeldir + "/best.pt"
-                opt.weights = self.weights
-        
-
-    def download_dataset_artifact(self,path, alias):
-        if path.startswith(WANDB_ARTIFACT_PREFIX):
-            dataset_artifact = wandb.use_artifact(remove_prefix(path,WANDB_ARTIFACT_PREFIX)+":"+alias)
-            if dataset_artifact is None:
-                logger.error('Error: W&B dataset artifact doesn\'t exist')
-                raise ValueError('Artifact doesn\'t exist')
-            datadir = dataset_artifact.download()
-            labels_zip = datadir+"/data/labels.zip"
-            shutil.unpack_archive(labels_zip, datadir+'/data/labels', 'zip')
-            print("Downloaded dataset to : ", datadir)
-            return datadir, dataset_artifact
-        return None, None
-    
-    def download_model_artifact(self,name):
-        model_artifact = wandb.use_artifact(name+":latest")
-        if model_artifact is None:
-            logger.error('Error: W&B model artifact doesn\'t exist')
-            raise ValueError('Artifact doesn\'t exist')
-        modeldir = model_artifact.download()
-        print("Downloaded model to : ", modeldir)
-        return modeldir, model_artifact
-    
-    def log_model(self, path, opt, epoch):
-        datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
-        model_artifact = wandb.Artifact('run_'+wandb.run.id+'_model', type='model', metadata={
-            'original_url': str(path),
-            'epoch': epoch+1,
-            'save period': opt.save_period,
-            'project': opt.project,
-            'datetime': datetime_suffix
-        })
-        model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
-        model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
-        wandb.log_artifact(model_artifact)
-
-        if epoch+1 == opt.epochs:
-            model_artifact = wandb.Artifact('final_model', type='model', metadata={
-            'run_id': wandb.run.id,
-            'datetime': datetime_suffix
-            })
-            model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
-            model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
-            wandb.log_artifact(model_artifact)
-        print("Saving model artifact on epoch ", epoch+1)
-    
-    def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
-        artifact = wandb.Artifact(name=name, type="dataset")
-        image_path = dataloader.dataset.path
-        artifact.add_dir(image_path,name='data/images')
-        table = wandb.Table(
-            columns=["id", "train_image", "Classes"]
-        )
-        id_count = 0
-        class_set = wandb.Classes([{'id':id , 'name':name} for id,name in class_to_id.items()])
-        for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
-            targets = targets.to(device)
-            nb, _, height, width = img.shape  # batch size, channels, height, width
-            targets[:,2:] = (xywh2xyxy(targets[:,2:].view(-1, 4)))
-            for si, _ in enumerate(img):
-                height, width = shapes[si][0]
-                labels = targets[targets[:, 0] == si]
-                labels[:,2:] *=  torch.Tensor([width, height, width, height]).to(device)
-                labels = labels[:, 1:]
-                box_data = []
-                img_classes = {}
-                for cls, *xyxy in labels.tolist():
-                  class_id = int(cls)
-                  box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                  "class_id": class_id,
-                                  "box_caption": "%s" % (class_to_id[class_id]),
-                                  "scores": {"acc": 1},
-                                  "domain": "pixel"})
-                  img_classes[class_id] = class_to_id[class_id]
-                boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-                table.add_data(id_count,wandb.Image(paths[si], classes=class_set, boxes=boxes), json.dumps(img_classes))
-                id_count = id_count+1
-        artifact.add(table, name)
-        label_path = image_path.replace('images','labels')
-        # Workaround for: Unable to log empty txt files via artifacts
-        if not os.path.isfile(name+'_labels.zip'): # make_archive won't check if file exists
-            shutil.make_archive(name+'_labels', 'zip', label_path)
-        artifact.add_file(name+'_labels.zip', name='data/labels.zip')
-        wandb.log_artifact(artifact)
-        print("Saving data to W&B...")
-            
-    def log(self, log_dict):
-        if self.wandb_run:
-            for key, value in log_dict.items():
-                self.log_dict[key] = value
-                
-    def end_epoch(self):
-        if self.wandb_run and self.log_dict:
-            wandb.log(self.log_dict)
-        self.log_dict = {}
-
-    def finish_run(self):
-        if self.wandb_run:
-            if self.result_artifact:
-                print("Add Training Progress Artifact")
-                self.result_artifact.add(self.result_table, 'result')
-                train_results = wandb.JoinedTable(self.testset_artifact.get("val"), self.result_table, "id")
-                self.result_artifact.add(train_results, 'joined_result')
-                wandb.log_artifact(self.result_artifact)
-            if self.log_dict:
-                wandb.log(self.log_dict)
-            wandb.run.finish()
diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 31633af635d0..bacc025e0d55 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,54 +1,66 @@
 import argparse
+import logging
 import os
+import random
+import time
+from pathlib import Path
+from threading import Thread
+from warnings import warn
 
+import numpy as np
 import torch
 import yaml
+import wandb
+from wandb_utils import WandbLogger
 
-from utils.datasets import create_dataloader
 from utils.general import check_dataset, colorstr
 from utils.torch_utils import torch_distributed_zero_first
-from wandb_utils import WandbLogger
+from utils.datasets import create_dataloader
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
-
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
-        data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
+        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
     wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
     # Hyperparameters
     with open(opt.hyp) as f:
-        hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
+        hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
+
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check
-
+    train_path = data_dict['train']
+    test_path = data_dict['val']
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
     imgsz, batch_size = opt.img_size, opt.batch_size
-    trainloader, trainset = create_dataloader(data_dict['train'], imgsz, batch_size, stride=32, opt=opt,
-                                              hyp=hyp, cache=False, rect=opt.rect, rank=-1,
-                                              world_size=opt.world_size, workers=opt.workers,
-                                              image_weights=opt.image_weights, quad=opt.quad,
-                                              prefix=colorstr('train: '))
-
-    testloader = create_dataloader(data_dict['val'], imgsz, batch_size, stride=32, opt=opt,  # testloader
-                                   hyp=hyp, cache=False, rect=True, rank=-1,
-                                   world_size=opt.world_size, workers=opt.workers,
-                                   pad=0.5, prefix=colorstr('val: '))[0]
+    assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
+    trainloader, trainset = create_dataloader(train_path, imgsz, batch_size, stride=32, opt=opt,
+                                            hyp=hyp, cache=False, rect=opt.rect, rank=-1,
+                                            world_size=opt.world_size, workers=opt.workers,
+                                            image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
+    mlc = np.concatenate(trainset.labels, 0)[:, 0].max()  # max label class
+    assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
+    
+    testloader = create_dataloader(test_path, imgsz, batch_size, stride=32, opt=opt,  # testloader
+                                       hyp=hyp, cache=False, rect=True, rank=-1,
+                                       world_size=opt.world_size, workers=opt.workers,
+                                       pad=0.5, prefix=colorstr('val: '))[0]
     names_to_ids = {k: v for k, v in enumerate(names)}
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
     wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
-    # Update/Create new config file with links to artifact
+    #Update/Create new config file with links to artifact
     data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
     data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
-    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')
-    data_dict.pop('download', None)  # Don't download the original dataset. Use artifacts
+    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.','_wandb.')
+    data_dict.pop('download',None) #Don't download the original dataset. Use artifacts
     with open(ouput_data_config, 'w') as fp:
         yaml.dump(data_dict, fp)
     print("New Config file => ", ouput_data_config)
-
-
+    
+    
+    
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index eac23ef9e841..85b3439df299 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -1,18 +1,18 @@
-import json
-import logging
-import os
-import shutil
-import sys
 from datetime import datetime
-from os.path import dirname
+import shutil
 from pathlib import Path
+import os
+from os.path import dirname
+import stat
+import logging
 
+import tqdm
 import torch
-
-# Add utils/ to path
+import json
+import sys
+# Add utils/ to path 
 sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))
 from utils.general import colorstr, xywh2xyxy
-
 logger = logging.getLogger(__name__)
 
 try:
@@ -23,20 +23,18 @@
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
-
 def remove_prefix(from_string, prefix):
     return from_string[len(prefix):]
-
-
+    
 class WandbLogger():
     def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
         self.wandb = wandb
         if self.wandb:
             self.wandb_run = wandb.init(config=opt, resume="allow",
-                                        project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
-                                        name=name,
-                                        job_type=job_type,
-                                        id=run_id)
+                                   project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                   name=name,
+                                   job_type=job_type,
+                                   id=run_id)
         else:
             self.wandb_run = None
         if job_type == 'Training':
@@ -45,13 +43,11 @@ def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
                 opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
             if opt.save_period == -1:
                 opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
-
+    
     def setup_training(self, opt, data_dict):
         self.log_dict = {}
-        self.train_artifact_path, self.trainset_artifact = \
-            self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
-        self.test_artifact_path, self.testset_artifact = \
-            self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
+        self.train_artifact_path, self.trainset_artifact = self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
+        self.test_artifact_path, self.testset_artifact = self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
         self.result_artifact, self.result_table, self.weights = None, None, None
         if self.train_artifact_path is not None:
             train_path = self.train_artifact_path + '/data/images/'
@@ -59,41 +55,42 @@ def setup_training(self, opt, data_dict):
         if self.test_artifact_path is not None:
             test_path = self.test_artifact_path + '/data/images/'
             data_dict['val'] = test_path
-            self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
+            self.result_artifact = wandb.Artifact("run_"+wandb.run.id+"_progress", "evaluation")
             self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
         if opt.resume_from_artifact:
             modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
             if modeldir:
                 self.weights = modeldir + "/best.pt"
                 opt.weights = self.weights
+        
 
-    def download_dataset_artifact(self, path, alias):
+    def download_dataset_artifact(self,path, alias):
         if path.startswith(WANDB_ARTIFACT_PREFIX):
-            dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
+            dataset_artifact = wandb.use_artifact(remove_prefix(path,WANDB_ARTIFACT_PREFIX)+":"+alias)
             if dataset_artifact is None:
                 logger.error('Error: W&B dataset artifact doesn\'t exist')
                 raise ValueError('Artifact doesn\'t exist')
             datadir = dataset_artifact.download()
-            labels_zip = datadir + "/data/labels.zip"
-            shutil.unpack_archive(labels_zip, datadir + '/data/labels', 'zip')
+            labels_zip = datadir+"/data/labels.zip"
+            shutil.unpack_archive(labels_zip, datadir+'/data/labels', 'zip')
             print("Downloaded dataset to : ", datadir)
             return datadir, dataset_artifact
         return None, None
-
-    def download_model_artifact(self, name):
-        model_artifact = wandb.use_artifact(name + ":latest")
+    
+    def download_model_artifact(self,name):
+        model_artifact = wandb.use_artifact(name+":latest")
         if model_artifact is None:
             logger.error('Error: W&B model artifact doesn\'t exist')
             raise ValueError('Artifact doesn\'t exist')
         modeldir = model_artifact.download()
         print("Downloaded model to : ", modeldir)
         return modeldir, model_artifact
-
+    
     def log_model(self, path, opt, epoch):
         datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
-        model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
+        model_artifact = wandb.Artifact('run_'+wandb.run.id+'_model', type='model', metadata={
             'original_url': str(path),
-            'epoch': epoch + 1,
+            'epoch': epoch+1,
             'save period': opt.save_period,
             'project': opt.project,
             'datetime': datetime_suffix
@@ -102,62 +99,61 @@ def log_model(self, path, opt, epoch):
         model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
         wandb.log_artifact(model_artifact)
 
-        if epoch + 1 == opt.epochs:
+        if epoch+1 == opt.epochs:
             model_artifact = wandb.Artifact('final_model', type='model', metadata={
-                'run_id': wandb.run.id,
-                'datetime': datetime_suffix
+            'run_id': wandb.run.id,
+            'datetime': datetime_suffix
             })
             model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
             model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
             wandb.log_artifact(model_artifact)
-        print("Saving model artifact on epoch ", epoch + 1)
-
+        print("Saving model artifact on epoch ", epoch+1)
+    
     def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
         image_path = dataloader.dataset.path
-        artifact.add_dir(image_path, name='data/images')
+        artifact.add_dir(image_path,name='data/images')
         table = wandb.Table(
             columns=["id", "train_image", "Classes"]
         )
         id_count = 0
-        class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
+        class_set = wandb.Classes([{'id':id , 'name':name} for id,name in class_to_id.items()])
         for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
             targets = targets.to(device)
             nb, _, height, width = img.shape  # batch size, channels, height, width
-            targets[:, 2:] = (xywh2xyxy(targets[:, 2:].view(-1, 4)))
+            targets[:,2:] = (xywh2xyxy(targets[:,2:].view(-1, 4)))
             for si, _ in enumerate(img):
                 height, width = shapes[si][0]
                 labels = targets[targets[:, 0] == si]
-                labels[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)
+                labels[:,2:] *=  torch.Tensor([width, height, width, height]).to(device)
                 labels = labels[:, 1:]
                 box_data = []
                 img_classes = {}
                 for cls, *xyxy in labels.tolist():
-                    class_id = int(cls)
-                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                     "class_id": class_id,
-                                     "box_caption": "%s" % (class_to_id[class_id]),
-                                     "scores": {"acc": 1},
-                                     "domain": "pixel"})
-                    img_classes[class_id] = class_to_id[class_id]
+                  class_id = int(cls)
+                  box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                  "class_id": class_id,
+                                  "box_caption": "%s" % (class_to_id[class_id]),
+                                  "scores": {"acc": 1},
+                                  "domain": "pixel"})
+                  img_classes[class_id] = class_to_id[class_id]
                 boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-                table.add_data(id_count, wandb.Image(paths[si], classes=class_set, boxes=boxes),
-                               json.dumps(img_classes))
-                id_count = id_count + 1
+                table.add_data(id_count,wandb.Image(paths[si], classes=class_set, boxes=boxes), json.dumps(img_classes))
+                id_count = id_count+1
         artifact.add(table, name)
-        label_path = image_path.replace('images', 'labels')
+        label_path = image_path.replace('images','labels')
         # Workaround for: Unable to log empty txt files via artifacts
-        if not os.path.isfile(name + '_labels.zip'):  # make_archive won't check if file exists
-            shutil.make_archive(name + '_labels', 'zip', label_path)
-        artifact.add_file(name + '_labels.zip', name='data/labels.zip')
+        if not os.path.isfile(name+'_labels.zip'): # make_archive won't check if file exists
+            shutil.make_archive(name+'_labels', 'zip', label_path)
+        artifact.add_file(name+'_labels.zip', name='data/labels.zip')
         wandb.log_artifact(artifact)
         print("Saving data to W&B...")
-
+            
     def log(self, log_dict):
         if self.wandb_run:
             for key, value in log_dict.items():
                 self.log_dict[key] = value
-
+                
     def end_epoch(self):
         if self.wandb_run and self.log_dict:
             wandb.log(self.log_dict)

From bbc5271c5de340087070d847b152160e171ebb78 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 20 Jan 2021 13:14:17 -0800
Subject: [PATCH 16/36] yaml.SafeLoader

---
 utils/wandb_logging/log_dataset.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index bacc025e0d55..a8e3eb9a7b54 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -21,11 +21,11 @@
 
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
-        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
+        data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
     wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
     # Hyperparameters
     with open(opt.hyp) as f:
-        hyp = yaml.load(f, Loader=yaml.FullLoader)  # load hyps
+        hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
 
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check

From 9bbecc870ecaa29ce42dca413ca70540d6241fa1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 20 Jan 2021 13:22:25 -0800
Subject: [PATCH 17/36] PEP8 reformat

---
 utils/wandb_logging/log_dataset.py |  38 +++++-----
 utils/wandb_logging/wandb_utils.py | 110 +++++++++++++++--------------
 2 files changed, 73 insertions(+), 75 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index a8e3eb9a7b54..f609fbe6d80f 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,24 +1,18 @@
 import argparse
-import logging
 import os
-import random
-import time
-from pathlib import Path
-from threading import Thread
-from warnings import warn
 
 import numpy as np
 import torch
 import yaml
-import wandb
-from wandb_utils import WandbLogger
 
+from utils.datasets import create_dataloader
 from utils.general import check_dataset, colorstr
 from utils.torch_utils import torch_distributed_zero_first
-from utils.datasets import create_dataloader
+from wandb_utils import WandbLogger
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
+
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
@@ -35,32 +29,32 @@ def create_dataset_artifact(opt):
     imgsz, batch_size = opt.img_size, opt.batch_size
     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
     trainloader, trainset = create_dataloader(train_path, imgsz, batch_size, stride=32, opt=opt,
-                                            hyp=hyp, cache=False, rect=opt.rect, rank=-1,
-                                            world_size=opt.world_size, workers=opt.workers,
-                                            image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
+                                              hyp=hyp, cache=False, rect=opt.rect, rank=-1,
+                                              world_size=opt.world_size, workers=opt.workers,
+                                              image_weights=opt.image_weights, quad=opt.quad,
+                                              prefix=colorstr('train: '))
     mlc = np.concatenate(trainset.labels, 0)[:, 0].max()  # max label class
     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
-    
+
     testloader = create_dataloader(test_path, imgsz, batch_size, stride=32, opt=opt,  # testloader
-                                       hyp=hyp, cache=False, rect=True, rank=-1,
-                                       world_size=opt.world_size, workers=opt.workers,
-                                       pad=0.5, prefix=colorstr('val: '))[0]
+                                   hyp=hyp, cache=False, rect=True, rank=-1,
+                                   world_size=opt.world_size, workers=opt.workers,
+                                   pad=0.5, prefix=colorstr('val: '))[0]
     names_to_ids = {k: v for k, v in enumerate(names)}
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
     wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
-    #Update/Create new config file with links to artifact
+    # Update/Create new config file with links to artifact
     data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
     data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
-    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.','_wandb.')
-    data_dict.pop('download',None) #Don't download the original dataset. Use artifacts
+    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')
+    data_dict.pop('download', None)  # Don't download the original dataset. Use artifacts
     with open(ouput_data_config, 'w') as fp:
         yaml.dump(data_dict, fp)
     print("New Config file => ", ouput_data_config)
-    
-    
-    
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 85b3439df299..eac23ef9e841 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -1,18 +1,18 @@
-from datetime import datetime
-import shutil
-from pathlib import Path
+import json
+import logging
 import os
+import shutil
+import sys
+from datetime import datetime
 from os.path import dirname
-import stat
-import logging
+from pathlib import Path
 
-import tqdm
 import torch
-import json
-import sys
-# Add utils/ to path 
+
+# Add utils/ to path
 sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))
 from utils.general import colorstr, xywh2xyxy
+
 logger = logging.getLogger(__name__)
 
 try:
@@ -23,18 +23,20 @@
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
+
 def remove_prefix(from_string, prefix):
     return from_string[len(prefix):]
-    
+
+
 class WandbLogger():
     def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
         self.wandb = wandb
         if self.wandb:
             self.wandb_run = wandb.init(config=opt, resume="allow",
-                                   project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
-                                   name=name,
-                                   job_type=job_type,
-                                   id=run_id)
+                                        project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                        name=name,
+                                        job_type=job_type,
+                                        id=run_id)
         else:
             self.wandb_run = None
         if job_type == 'Training':
@@ -43,11 +45,13 @@ def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
                 opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
             if opt.save_period == -1:
                 opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
-    
+
     def setup_training(self, opt, data_dict):
         self.log_dict = {}
-        self.train_artifact_path, self.trainset_artifact = self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
-        self.test_artifact_path, self.testset_artifact = self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
+        self.train_artifact_path, self.trainset_artifact = \
+            self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
+        self.test_artifact_path, self.testset_artifact = \
+            self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
         self.result_artifact, self.result_table, self.weights = None, None, None
         if self.train_artifact_path is not None:
             train_path = self.train_artifact_path + '/data/images/'
@@ -55,42 +59,41 @@ def setup_training(self, opt, data_dict):
         if self.test_artifact_path is not None:
             test_path = self.test_artifact_path + '/data/images/'
             data_dict['val'] = test_path
-            self.result_artifact = wandb.Artifact("run_"+wandb.run.id+"_progress", "evaluation")
+            self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
             self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
         if opt.resume_from_artifact:
             modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
             if modeldir:
                 self.weights = modeldir + "/best.pt"
                 opt.weights = self.weights
-        
 
-    def download_dataset_artifact(self,path, alias):
+    def download_dataset_artifact(self, path, alias):
         if path.startswith(WANDB_ARTIFACT_PREFIX):
-            dataset_artifact = wandb.use_artifact(remove_prefix(path,WANDB_ARTIFACT_PREFIX)+":"+alias)
+            dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
             if dataset_artifact is None:
                 logger.error('Error: W&B dataset artifact doesn\'t exist')
                 raise ValueError('Artifact doesn\'t exist')
             datadir = dataset_artifact.download()
-            labels_zip = datadir+"/data/labels.zip"
-            shutil.unpack_archive(labels_zip, datadir+'/data/labels', 'zip')
+            labels_zip = datadir + "/data/labels.zip"
+            shutil.unpack_archive(labels_zip, datadir + '/data/labels', 'zip')
             print("Downloaded dataset to : ", datadir)
             return datadir, dataset_artifact
         return None, None
-    
-    def download_model_artifact(self,name):
-        model_artifact = wandb.use_artifact(name+":latest")
+
+    def download_model_artifact(self, name):
+        model_artifact = wandb.use_artifact(name + ":latest")
         if model_artifact is None:
             logger.error('Error: W&B model artifact doesn\'t exist')
             raise ValueError('Artifact doesn\'t exist')
         modeldir = model_artifact.download()
         print("Downloaded model to : ", modeldir)
         return modeldir, model_artifact
-    
+
     def log_model(self, path, opt, epoch):
         datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
-        model_artifact = wandb.Artifact('run_'+wandb.run.id+'_model', type='model', metadata={
+        model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
             'original_url': str(path),
-            'epoch': epoch+1,
+            'epoch': epoch + 1,
             'save period': opt.save_period,
             'project': opt.project,
             'datetime': datetime_suffix
@@ -99,61 +102,62 @@ def log_model(self, path, opt, epoch):
         model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
         wandb.log_artifact(model_artifact)
 
-        if epoch+1 == opt.epochs:
+        if epoch + 1 == opt.epochs:
             model_artifact = wandb.Artifact('final_model', type='model', metadata={
-            'run_id': wandb.run.id,
-            'datetime': datetime_suffix
+                'run_id': wandb.run.id,
+                'datetime': datetime_suffix
             })
             model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
             model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
             wandb.log_artifact(model_artifact)
-        print("Saving model artifact on epoch ", epoch+1)
-    
+        print("Saving model artifact on epoch ", epoch + 1)
+
     def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
         image_path = dataloader.dataset.path
-        artifact.add_dir(image_path,name='data/images')
+        artifact.add_dir(image_path, name='data/images')
         table = wandb.Table(
             columns=["id", "train_image", "Classes"]
         )
         id_count = 0
-        class_set = wandb.Classes([{'id':id , 'name':name} for id,name in class_to_id.items()])
+        class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
         for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
             targets = targets.to(device)
             nb, _, height, width = img.shape  # batch size, channels, height, width
-            targets[:,2:] = (xywh2xyxy(targets[:,2:].view(-1, 4)))
+            targets[:, 2:] = (xywh2xyxy(targets[:, 2:].view(-1, 4)))
             for si, _ in enumerate(img):
                 height, width = shapes[si][0]
                 labels = targets[targets[:, 0] == si]
-                labels[:,2:] *=  torch.Tensor([width, height, width, height]).to(device)
+                labels[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)
                 labels = labels[:, 1:]
                 box_data = []
                 img_classes = {}
                 for cls, *xyxy in labels.tolist():
-                  class_id = int(cls)
-                  box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                  "class_id": class_id,
-                                  "box_caption": "%s" % (class_to_id[class_id]),
-                                  "scores": {"acc": 1},
-                                  "domain": "pixel"})
-                  img_classes[class_id] = class_to_id[class_id]
+                    class_id = int(cls)
+                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                     "class_id": class_id,
+                                     "box_caption": "%s" % (class_to_id[class_id]),
+                                     "scores": {"acc": 1},
+                                     "domain": "pixel"})
+                    img_classes[class_id] = class_to_id[class_id]
                 boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-                table.add_data(id_count,wandb.Image(paths[si], classes=class_set, boxes=boxes), json.dumps(img_classes))
-                id_count = id_count+1
+                table.add_data(id_count, wandb.Image(paths[si], classes=class_set, boxes=boxes),
+                               json.dumps(img_classes))
+                id_count = id_count + 1
         artifact.add(table, name)
-        label_path = image_path.replace('images','labels')
+        label_path = image_path.replace('images', 'labels')
         # Workaround for: Unable to log empty txt files via artifacts
-        if not os.path.isfile(name+'_labels.zip'): # make_archive won't check if file exists
-            shutil.make_archive(name+'_labels', 'zip', label_path)
-        artifact.add_file(name+'_labels.zip', name='data/labels.zip')
+        if not os.path.isfile(name + '_labels.zip'):  # make_archive won't check if file exists
+            shutil.make_archive(name + '_labels', 'zip', label_path)
+        artifact.add_file(name + '_labels.zip', name='data/labels.zip')
         wandb.log_artifact(artifact)
         print("Saving data to W&B...")
-            
+
     def log(self, log_dict):
         if self.wandb_run:
             for key, value in log_dict.items():
                 self.log_dict[key] = value
-                
+
     def end_epoch(self):
         if self.wandb_run and self.log_dict:
             wandb.log(self.log_dict)

From ceb63cd6a19532a728b9e0c8e189d87754322d47 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 20 Jan 2021 13:28:29 -0800
Subject: [PATCH 18/36] remove redundant checks

---
 utils/wandb_logging/log_dataset.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index f609fbe6d80f..31633af635d0 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,7 +1,6 @@
 import argparse
 import os
 
-import numpy as np
 import torch
 import yaml
 
@@ -20,23 +19,18 @@ def create_dataset_artifact(opt):
     # Hyperparameters
     with open(opt.hyp) as f:
         hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
-
     with torch_distributed_zero_first(-1):
         check_dataset(data_dict)  # check
-    train_path = data_dict['train']
-    test_path = data_dict['val']
+
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
     imgsz, batch_size = opt.img_size, opt.batch_size
-    assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
-    trainloader, trainset = create_dataloader(train_path, imgsz, batch_size, stride=32, opt=opt,
+    trainloader, trainset = create_dataloader(data_dict['train'], imgsz, batch_size, stride=32, opt=opt,
                                               hyp=hyp, cache=False, rect=opt.rect, rank=-1,
                                               world_size=opt.world_size, workers=opt.workers,
                                               image_weights=opt.image_weights, quad=opt.quad,
                                               prefix=colorstr('train: '))
-    mlc = np.concatenate(trainset.labels, 0)[:, 0].max()  # max label class
-    assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
 
-    testloader = create_dataloader(test_path, imgsz, batch_size, stride=32, opt=opt,  # testloader
+    testloader = create_dataloader(data_dict['val'], imgsz, batch_size, stride=32, opt=opt,  # testloader
                                    hyp=hyp, cache=False, rect=True, rank=-1,
                                    world_size=opt.world_size, workers=opt.workers,
                                    pad=0.5, prefix=colorstr('val: '))[0]

From 2529fb791e643c6d9a91c83315b513825af296de Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Thu, 21 Jan 2021 06:59:17 +0000
Subject: [PATCH 19/36] Update util files

---
 utils/wandb_logging/log_dataset.py | 40 +++++----------------
 utils/wandb_logging/wandb_utils.py | 58 ++++++++++++++----------------
 2 files changed, 36 insertions(+), 62 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 31633af635d0..45c800356b88 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,13 +1,14 @@
 import argparse
 import os
+from pathlib import Path
 
 import torch
 import yaml
 
-from utils.datasets import create_dataloader
+from wandb_utils import WandbLogger
+from utils.datasets import create_dataloader, LoadImagesAndLabels
 from utils.general import check_dataset, colorstr
 from utils.torch_utils import torch_distributed_zero_first
-from wandb_utils import WandbLogger
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 
@@ -16,32 +17,15 @@ def create_dataset_artifact(opt):
     with open(opt.data) as f:
         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
     wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
-    # Hyperparameters
-    with open(opt.hyp) as f:
-        hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
-    with torch_distributed_zero_first(-1):
-        check_dataset(data_dict)  # check
-
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
-    imgsz, batch_size = opt.img_size, opt.batch_size
-    trainloader, trainset = create_dataloader(data_dict['train'], imgsz, batch_size, stride=32, opt=opt,
-                                              hyp=hyp, cache=False, rect=opt.rect, rank=-1,
-                                              world_size=opt.world_size, workers=opt.workers,
-                                              image_weights=opt.image_weights, quad=opt.quad,
-                                              prefix=colorstr('train: '))
-
-    testloader = create_dataloader(data_dict['val'], imgsz, batch_size, stride=32, opt=opt,  # testloader
-                                   hyp=hyp, cache=False, rect=True, rank=-1,
-                                   world_size=opt.world_size, workers=opt.workers,
-                                   pad=0.5, prefix=colorstr('val: '))[0]
+    trainset = LoadImagesAndLabels(data_dict['train'], rect=opt.rect)
+    testset = LoadImagesAndLabels(data_dict['val'], rect=opt.rect)
     names_to_ids = {k: v for k, v in enumerate(names)}
-
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    wandb_logger.log_dataset_artifact(trainloader, device, names_to_ids, name='train')
-    wandb_logger.log_dataset_artifact(testloader, device, names_to_ids, name='val')
+    wandb_logger.log_dataset_artifact(trainset, names_to_ids, name='train')
+    wandb_logger.log_dataset_artifact(testset, names_to_ids, name='val')
     # Update/Create new config file with links to artifact
-    data_dict['train'] = WANDB_ARTIFACT_PREFIX + opt.project + '/train'
-    data_dict['val'] = WANDB_ARTIFACT_PREFIX + opt.project + '/val'
+    data_dict['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
+    data_dict['val'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'val')
     ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')
     data_dict.pop('download', None)  # Don't download the original dataset. Use artifacts
     with open(ouput_data_config, 'w') as fp:
@@ -53,14 +37,8 @@ def create_dataset_artifact(opt):
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
-    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
     parser.add_argument('--rect', action='store_true', help='rectangular training')
-    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
     parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
-    parser.add_argument('--img-size', nargs='+', type=int, default=640, help='[train, test] image sizes')
-    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
-    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
-    parser.add_argument('--quad', action='store_true', help='quad dataloader')
     parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
     opt = parser.parse_args()
     opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index eac23ef9e841..de211eb95e09 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -54,17 +54,17 @@ def setup_training(self, opt, data_dict):
             self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
         self.result_artifact, self.result_table, self.weights = None, None, None
         if self.train_artifact_path is not None:
-            train_path = self.train_artifact_path + '/data/images/'
+            train_path = Path(self.train_artifact_path) / '/data/images/'
             data_dict['train'] = train_path
         if self.test_artifact_path is not None:
-            test_path = self.test_artifact_path + '/data/images/'
+            test_path = Path(self.test_artifact_path) / '/data/images/'
             data_dict['val'] = test_path
             self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
             self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
         if opt.resume_from_artifact:
             modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
             if modeldir:
-                self.weights = modeldir + "/best.pt"
+                self.weights = Path(modeldir) / "/best.pt"
                 opt.weights = self.weights
 
     def download_dataset_artifact(self, path, alias):
@@ -74,8 +74,8 @@ def download_dataset_artifact(self, path, alias):
                 logger.error('Error: W&B dataset artifact doesn\'t exist')
                 raise ValueError('Artifact doesn\'t exist')
             datadir = dataset_artifact.download()
-            labels_zip = datadir + "/data/labels.zip"
-            shutil.unpack_archive(labels_zip, datadir + '/data/labels', 'zip')
+            labels_zip = Path(datadir) / "/data/labels.zip"
+            shutil.unpack_archive(labels_zip, Path(datadir) / '/data/labels', 'zip')
             print("Downloaded dataset to : ", datadir)
             return datadir, dataset_artifact
         return None, None
@@ -112,41 +112,37 @@ def log_model(self, path, opt, epoch):
             wandb.log_artifact(model_artifact)
         print("Saving model artifact on epoch ", epoch + 1)
 
-    def log_dataset_artifact(self, dataloader, device, class_to_id, name='dataset'):
+    def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
-        image_path = dataloader.dataset.path
+        image_path = dataloader.path
         artifact.add_dir(image_path, name='data/images')
         table = wandb.Table(
             columns=["id", "train_image", "Classes"]
         )
         id_count = 0
         class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
-        for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
-            targets = targets.to(device)
-            nb, _, height, width = img.shape  # batch size, channels, height, width
-            targets[:, 2:] = (xywh2xyxy(targets[:, 2:].view(-1, 4)))
-            for si, _ in enumerate(img):
-                height, width = shapes[si][0]
-                labels = targets[targets[:, 0] == si]
-                labels[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)
-                labels = labels[:, 1:]
-                box_data = []
-                img_classes = {}
-                for cls, *xyxy in labels.tolist():
-                    class_id = int(cls)
-                    box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                     "class_id": class_id,
-                                     "box_caption": "%s" % (class_to_id[class_id]),
-                                     "scores": {"acc": 1},
-                                     "domain": "pixel"})
-                    img_classes[class_id] = class_to_id[class_id]
-                boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-                table.add_data(id_count, wandb.Image(paths[si], classes=class_set, boxes=boxes),
-                               json.dumps(img_classes))
-                id_count = id_count + 1
+        for si, (img, labels, paths, shapes) in enumerate(dataloader):
+            _, height, width = img.shape  # batch size, channels, height, width
+            labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
+            height, width = shapes[0]
+            labels[:, 2:] *= torch.Tensor([width, height, width, height])
+            labels = labels[:, 1:]
+            box_data = []
+            img_classes = {}
+            for cls, *xyxy in labels.tolist():
+                class_id = int(cls)
+                box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                 "class_id": class_id,
+                                 "box_caption": "%s" % (class_to_id[class_id]),
+                                 "scores": {"acc": 1},
+                                 "domain": "pixel"})
+                img_classes[class_id] = class_to_id[class_id]
+            boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
+            table.add_data(id_count, wandb.Image(paths, classes=class_set, boxes=boxes),
+                           json.dumps(img_classes))
+            id_count = id_count + 1
         artifact.add(table, name)
         label_path = image_path.replace('images', 'labels')
-        # Workaround for: Unable to log empty txt files via artifacts
         if not os.path.isfile(name + '_labels.zip'):  # make_archive won't check if file exists
             shutil.make_archive(name + '_labels', 'zip', label_path)
         artifact.add_file(name + '_labels.zip', name='data/labels.zip')

From efdcd7ee5a8e97895994acc6c69fb91bfee6d220 Mon Sep 17 00:00:00 2001
From: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Thu, 21 Jan 2021 13:16:13 +0530
Subject: [PATCH 20/36] Update wandb_utils.py

---
 utils/wandb_logging/wandb_utils.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index de211eb95e09..77baa17c4025 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -54,17 +54,17 @@ def setup_training(self, opt, data_dict):
             self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
         self.result_artifact, self.result_table, self.weights = None, None, None
         if self.train_artifact_path is not None:
-            train_path = Path(self.train_artifact_path) / '/data/images/'
-            data_dict['train'] = train_path
+            train_path = Path(self.train_artifact_path) / 'data/images/'
+            data_dict['train'] = str(train_path)
         if self.test_artifact_path is not None:
-            test_path = Path(self.test_artifact_path) / '/data/images/'
-            data_dict['val'] = test_path
+            test_path = Path(self.test_artifact_path) / 'data/images/'
+            data_dict['val'] = str(test_path)
             self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
             self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
         if opt.resume_from_artifact:
             modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
             if modeldir:
-                self.weights = Path(modeldir) / "/best.pt"
+                self.weights = Path(modeldir) / "best.pt"
                 opt.weights = self.weights
 
     def download_dataset_artifact(self, path, alias):
@@ -74,8 +74,8 @@ def download_dataset_artifact(self, path, alias):
                 logger.error('Error: W&B dataset artifact doesn\'t exist')
                 raise ValueError('Artifact doesn\'t exist')
             datadir = dataset_artifact.download()
-            labels_zip = Path(datadir) / "/data/labels.zip"
-            shutil.unpack_archive(labels_zip, Path(datadir) / '/data/labels', 'zip')
+            labels_zip = Path(datadir) / "data/labels.zip"
+            shutil.unpack_archive(labels_zip, Path(datadir) / 'data/labels', 'zip')
             print("Downloaded dataset to : ", datadir)
             return datadir, dataset_artifact
         return None, None

From 0d84870c72c6858c222da9c166c9e1493e4508ef Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Thu, 21 Jan 2021 13:08:48 +0000
Subject: [PATCH 21/36] Remove word size

---
 utils/wandb_logging/log_dataset.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 45c800356b88..26f957052f60 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -41,6 +41,5 @@ def create_dataset_artifact(opt):
     parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
     parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
     opt = parser.parse_args()
-    opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
 
     create_dataset_artifact(opt)

From 93219b9219f2fbd4c2b94ccc83360c1221e3bfac Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Thu, 21 Jan 2021 13:19:35 +0000
Subject: [PATCH 22/36] Change path of labels.zip

---
 utils/wandb_logging/wandb_utils.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 77baa17c4025..743f75862c9e 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -142,10 +142,11 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
                            json.dumps(img_classes))
             id_count = id_count + 1
         artifact.add(table, name)
-        label_path = image_path.replace('images', 'labels')
-        if not os.path.isfile(name + '_labels.zip'):  # make_archive won't check if file exists
-            shutil.make_archive(name + '_labels', 'zip', label_path)
-        artifact.add_file(name + '_labels.zip', name='data/labels.zip')
+        labels_path = image_path.replace('images', 'labels')
+        labels_zipped_path = Path(labels_path).parent / (name + '_labels.zip')
+        if not os.path.isfile(labels_zipped_path):  # make_archive won't check if file exists
+            shutil.make_archive(Path(labels_path).parent / (name + '_labels'), 'zip', labels_path)
+        artifact.add_file(str(labels_zipped_path), name='data/labels.zip')
         wandb.log_artifact(artifact)
         print("Saving data to W&B...")
 

From ac8fad283e05a89451709e440c51d2709aa3143f Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 13:32:58 -0800
Subject: [PATCH 23/36] remove unused imports

---
 utils/wandb_logging/log_dataset.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 26f957052f60..1db2689404c3 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -1,14 +1,10 @@
 import argparse
-import os
 from pathlib import Path
 
-import torch
 import yaml
 
+from utils.datasets import LoadImagesAndLabels
 from wandb_utils import WandbLogger
-from utils.datasets import create_dataloader, LoadImagesAndLabels
-from utils.general import check_dataset, colorstr
-from utils.torch_utils import torch_distributed_zero_first
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 

From e40b817dcfa8e4b739137f8d9271d7c516c5075f Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 13:41:57 -0800
Subject: [PATCH 24/36] remove --rect

---
 utils/wandb_logging/log_dataset.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 1db2689404c3..57359f94368d 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -14,8 +14,8 @@ def create_dataset_artifact(opt):
         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
     wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
-    trainset = LoadImagesAndLabels(data_dict['train'], rect=opt.rect)
-    testset = LoadImagesAndLabels(data_dict['val'], rect=opt.rect)
+    trainset = LoadImagesAndLabels(data_dict['train'])
+    testset = LoadImagesAndLabels(data_dict['val'])
     names_to_ids = {k: v for k, v in enumerate(names)}
     wandb_logger.log_dataset_artifact(trainset, names_to_ids, name='train')
     wandb_logger.log_dataset_artifact(testset, names_to_ids, name='val')
@@ -33,7 +33,6 @@ def create_dataset_artifact(opt):
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
-    parser.add_argument('--rect', action='store_true', help='rectangular training')
     parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
     parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
     opt = parser.parse_args()

From 0e219896eedef7673b2b7417f6edb8d0304ba7b1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 13:53:57 -0800
Subject: [PATCH 25/36] log_dataset.py cleanup

---
 utils/wandb_logging/log_dataset.py | 31 +++++++++++++++---------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 57359f94368d..ae6d3e81bcf5 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -11,22 +11,21 @@
 
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
-        data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
-    wandb_logger = WandbLogger(opt, '', None, data_dict, job_type='create_dataset')
-    nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])
-    trainset = LoadImagesAndLabels(data_dict['train'])
-    testset = LoadImagesAndLabels(data_dict['val'])
+        data = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
+    logger = WandbLogger(opt, '', None, data, job_type='create_dataset')
+    nc, names = (1, ['item']) if opt.single_cls else (int(data['nc']), data['names'])
     names_to_ids = {k: v for k, v in enumerate(names)}
-    wandb_logger.log_dataset_artifact(trainset, names_to_ids, name='train')
-    wandb_logger.log_dataset_artifact(testset, names_to_ids, name='val')
-    # Update/Create new config file with links to artifact
-    data_dict['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
-    data_dict['val'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'val')
-    ouput_data_config = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')
-    data_dict.pop('download', None)  # Don't download the original dataset. Use artifacts
-    with open(ouput_data_config, 'w') as fp:
-        yaml.dump(data_dict, fp)
-    print("New Config file => ", ouput_data_config)
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['train']), names_to_ids, name='train')  # trainset
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['val']), names_to_ids, name='val')  # valset
+
+    # Update data.yaml with artifact links
+    data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
+    data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'val')
+    path = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')  # updated data.yaml path
+    data.pop('download', None)  # download via artifact instead of predefined field 'download:'
+    with open(path, 'w') as f:
+        yaml.dump(data, f)
+    print("New Config file => ", path)
 
 
 if __name__ == '__main__':
@@ -34,7 +33,7 @@ def create_dataset_artifact(opt):
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
     parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
-    parser.add_argument('--overwrite_config', action='store_true', help='replace the origin data config file')
+    parser.add_argument('--overwrite_config', action='store_true', help='overwrite data.yaml')
     opt = parser.parse_args()
 
     create_dataset_artifact(opt)

From 3e538651cf4fc9b4861812efc8282cf2a0a1d392 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 13:57:49 -0800
Subject: [PATCH 26/36] log_dataset.py cleanup2

---
 utils/wandb_logging/log_dataset.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index ae6d3e81bcf5..72febc86f7d9 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -14,9 +14,9 @@ def create_dataset_artifact(opt):
         data = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
     logger = WandbLogger(opt, '', None, data, job_type='create_dataset')
     nc, names = (1, ['item']) if opt.single_cls else (int(data['nc']), data['names'])
-    names_to_ids = {k: v for k, v in enumerate(names)}
-    logger.log_dataset_artifact(LoadImagesAndLabels(data['train']), names_to_ids, name='train')  # trainset
-    logger.log_dataset_artifact(LoadImagesAndLabels(data['val']), names_to_ids, name='val')  # valset
+    names = {k: v for k, v in enumerate(names)}  # to index dictionary
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['train']), names, name='train')  # trainset
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['val']), names, name='val')  # valset
 
     # Update data.yaml with artifact links
     data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
@@ -32,7 +32,7 @@ def create_dataset_artifact(opt):
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
-    parser.add_argument('--project', type=str, default='yolov5', help='name of W&B Project')
+    parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
     parser.add_argument('--overwrite_config', action='store_true', help='overwrite data.yaml')
     opt = parser.parse_args()
 

From 4ebfb830972bf07ed947bcf6e156cd10b6d82f8a Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 14:10:12 -0800
Subject: [PATCH 27/36] wandb_utils.py cleanup

---
 utils/wandb_logging/wandb_utils.py | 21 +++++++++------------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 743f75862c9e..763077c85808 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -9,11 +9,9 @@
 
 import torch
 
-# Add utils/ to path
-sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))
-from utils.general import colorstr, xywh2xyxy
+sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))  # add utils/ to path
 
-logger = logging.getLogger(__name__)
+from utils.general import colorstr, xywh2xyxy
 
 try:
     import wandb
@@ -22,6 +20,7 @@
     print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+logger = logging.getLogger(__name__)
 
 
 def remove_prefix(from_string, prefix):
@@ -31,14 +30,12 @@ def remove_prefix(from_string, prefix):
 class WandbLogger():
     def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
         self.wandb = wandb
-        if self.wandb:
-            self.wandb_run = wandb.init(config=opt, resume="allow",
-                                        project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
-                                        name=name,
-                                        job_type=job_type,
-                                        id=run_id)
-        else:
-            self.wandb_run = None
+        self.wandb_run = wandb.init(config=opt, resume="allow",
+                                    project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                    name=name,
+                                    job_type=job_type,
+                                    id=run_id) if self.wandb else None
+
         if job_type == 'Training':
             self.setup_training(opt, data_dict)
             if opt.bbox_interval == -1:

From 58616a68e8f74fb164f60443af0f6de7ffc31981 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 14:11:37 -0800
Subject: [PATCH 28/36] remove redundant id_count

---
 utils/wandb_logging/wandb_utils.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 763077c85808..408fb75d31eb 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -113,10 +113,7 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
         image_path = dataloader.path
         artifact.add_dir(image_path, name='data/images')
-        table = wandb.Table(
-            columns=["id", "train_image", "Classes"]
-        )
-        id_count = 0
+        table = wandb.Table(columns=["id", "train_image", "Classes"])
         class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
         for si, (img, labels, paths, shapes) in enumerate(dataloader):
             _, height, width = img.shape  # batch size, channels, height, width
@@ -135,9 +132,8 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
                                  "domain": "pixel"})
                 img_classes[class_id] = class_to_id[class_id]
             boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-            table.add_data(id_count, wandb.Image(paths, classes=class_set, boxes=boxes),
+            table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes),
                            json.dumps(img_classes))
-            id_count = id_count + 1
         artifact.add(table, name)
         labels_path = image_path.replace('images', 'labels')
         labels_zipped_path = Path(labels_path).parent / (name + '_labels.zip')

From a8168c3be6a29bd8271294b15fb27a20a12e2b7a Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 14:15:22 -0800
Subject: [PATCH 29/36] wandb_utils.py cleanup2

---
 utils/wandb_logging/wandb_utils.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 408fb75d31eb..f3518f2d115d 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -132,8 +132,7 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
                                  "domain": "pixel"})
                 img_classes[class_id] = class_to_id[class_id]
             boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-            table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes),
-                           json.dumps(img_classes))
+            table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes))
         artifact.add(table, name)
         labels_path = image_path.replace('images', 'labels')
         labels_zipped_path = Path(labels_path).parent / (name + '_labels.zip')

From 1bb261d59d4736878e764ee9dff5a845ee7ec679 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 14:31:50 -0800
Subject: [PATCH 30/36] rename cls

---
 utils/wandb_logging/wandb_utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index f3518f2d115d..d2998c85b881 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -124,13 +124,13 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
             box_data = []
             img_classes = {}
             for cls, *xyxy in labels.tolist():
-                class_id = int(cls)
+                cls = int(cls)
                 box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
-                                 "class_id": class_id,
-                                 "box_caption": "%s" % (class_to_id[class_id]),
+                                 "class_id": cls,
+                                 "box_caption": "%s" % (class_to_id[cls]),
                                  "scores": {"acc": 1},
                                  "domain": "pixel"})
-                img_classes[class_id] = class_to_id[class_id]
+                img_classes[cls] = class_to_id[cls]
             boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
             table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes))
         artifact.add(table, name)

From 850b59da1a19409a9d17ad89ca66f3aab92a719f Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 15:04:39 -0800
Subject: [PATCH 31/36] use pathlib for zip

---
 utils/wandb_logging/wandb_utils.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index d2998c85b881..d2ee7f67daa8 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -120,10 +120,9 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
             labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
             height, width = shapes[0]
             labels[:, 2:] *= torch.Tensor([width, height, width, height])
-            labels = labels[:, 1:]
             box_data = []
             img_classes = {}
-            for cls, *xyxy in labels.tolist():
+            for cls, *xyxy in labels[:, 1:].tolist():
                 cls = int(cls)
                 box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
                                  "class_id": cls,
@@ -135,10 +134,10 @@ def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
             table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes))
         artifact.add(table, name)
         labels_path = image_path.replace('images', 'labels')
-        labels_zipped_path = Path(labels_path).parent / (name + '_labels.zip')
-        if not os.path.isfile(labels_zipped_path):  # make_archive won't check if file exists
-            shutil.make_archive(Path(labels_path).parent / (name + '_labels'), 'zip', labels_path)
-        artifact.add_file(str(labels_zipped_path), name='data/labels.zip')
+        zip_path = Path(labels_path).parent / (name + '_labels.zip')
+        if not zip_path.is_file():  # make_archive won't check if file exists
+            shutil.make_archive(zip_path.with_suffix(''), 'zip', labels_path)
+        artifact.add_file(str(zip_path), name='data/labels.zip')
         wandb.log_artifact(artifact)
         print("Saving data to W&B...")
 

From 94b7631739c766827f480f612c51a053f2d131d1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 23 Jan 2021 15:12:34 -0800
Subject: [PATCH 32/36] rename dataloader to dataset

---
 utils/wandb_logging/wandb_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index d2ee7f67daa8..01dd2f5a1b00 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -109,13 +109,13 @@ def log_model(self, path, opt, epoch):
             wandb.log_artifact(model_artifact)
         print("Saving model artifact on epoch ", epoch + 1)
 
-    def log_dataset_artifact(self, dataloader, class_to_id, name='dataset'):
+    def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
-        image_path = dataloader.path
+        image_path = dataset.path
         artifact.add_dir(image_path, name='data/images')
         table = wandb.Table(columns=["id", "train_image", "Classes"])
         class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
-        for si, (img, labels, paths, shapes) in enumerate(dataloader):
+        for si, (img, labels, paths, shapes) in enumerate(dataset):
             _, height, width = img.shape  # batch size, channels, height, width
             labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
             height, width = shapes[0]

From b959a02f770a63491c1168e098eebe380a67bd79 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Wed, 27 Jan 2021 05:54:35 +0000
Subject: [PATCH 33/36] Change import order

---
 utils/wandb_logging/log_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/wandb_logging/log_dataset.py b/utils/wandb_logging/log_dataset.py
index 72febc86f7d9..d790a9ce721e 100644
--- a/utils/wandb_logging/log_dataset.py
+++ b/utils/wandb_logging/log_dataset.py
@@ -3,8 +3,8 @@
 
 import yaml
 
-from utils.datasets import LoadImagesAndLabels
 from wandb_utils import WandbLogger
+from utils.datasets import LoadImagesAndLabels
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 

From 4e9a54c157170b34a890a7e324a0905f3d79c84d Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Sat, 30 Jan 2021 07:43:10 +0000
Subject: [PATCH 34/36] Remove redundant code

---
 utils/wandb_logging/wandb_utils.py | 29 ++++++-----------------------
 1 file changed, 6 insertions(+), 23 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 01dd2f5a1b00..7476000196de 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -4,15 +4,13 @@
 import shutil
 import sys
 from datetime import datetime
-from os.path import dirname
 from pathlib import Path
 
 import torch
 
-sys.path.append(dirname(dirname(dirname(os.path.abspath(__file__)))))  # add utils/ to path
-
+sys.path.append(str(Path(__file__).parent.parent.parent))  # add utils/ to path
 from utils.general import colorstr, xywh2xyxy
-
+from utils.datasets import img2label_paths
 try:
     import wandb
 except ImportError:
@@ -20,7 +18,6 @@
     print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
 
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
-logger = logging.getLogger(__name__)
 
 
 def remove_prefix(from_string, prefix):
@@ -67,9 +64,7 @@ def setup_training(self, opt, data_dict):
     def download_dataset_artifact(self, path, alias):
         if path.startswith(WANDB_ARTIFACT_PREFIX):
             dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
-            if dataset_artifact is None:
-                logger.error('Error: W&B dataset artifact doesn\'t exist')
-                raise ValueError('Artifact doesn\'t exist')
+            assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
             datadir = dataset_artifact.download()
             labels_zip = Path(datadir) / "data/labels.zip"
             shutil.unpack_archive(labels_zip, Path(datadir) / 'data/labels', 'zip')
@@ -79,9 +74,7 @@ def download_dataset_artifact(self, path, alias):
 
     def download_model_artifact(self, name):
         model_artifact = wandb.use_artifact(name + ":latest")
-        if model_artifact is None:
-            logger.error('Error: W&B model artifact doesn\'t exist')
-            raise ValueError('Artifact doesn\'t exist')
+        assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
         modeldir = model_artifact.download()
         print("Downloaded model to : ", modeldir)
         return modeldir, model_artifact
@@ -98,15 +91,6 @@ def log_model(self, path, opt, epoch):
         model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
         model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
         wandb.log_artifact(model_artifact)
-
-        if epoch + 1 == opt.epochs:
-            model_artifact = wandb.Artifact('final_model', type='model', metadata={
-                'run_id': wandb.run.id,
-                'datetime': datetime_suffix
-            })
-            model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
-            model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
-            wandb.log_artifact(model_artifact)
         print("Saving model artifact on epoch ", epoch + 1)
 
     def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
@@ -116,9 +100,8 @@ def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
         table = wandb.Table(columns=["id", "train_image", "Classes"])
         class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
         for si, (img, labels, paths, shapes) in enumerate(dataset):
-            _, height, width = img.shape  # batch size, channels, height, width
-            labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
             height, width = shapes[0]
+            labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
             labels[:, 2:] *= torch.Tensor([width, height, width, height])
             box_data = []
             img_classes = {}
@@ -133,7 +116,7 @@ def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
             boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
             table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes))
         artifact.add(table, name)
-        labels_path = image_path.replace('images', 'labels')
+        labels_path = 'labels'.join(image_path.rsplit('images', 1))
         zip_path = Path(labels_path).parent / (name + '_labels.zip')
         if not zip_path.is_file():  # make_archive won't check if file exists
             shutil.make_archive(zip_path.with_suffix(''), 'zip', labels_path)

From 28ad4ab2bccb2f7b9c0cff2c68035cadd4358309 Mon Sep 17 00:00:00 2001
From: AYush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Sat, 30 Jan 2021 07:48:26 +0000
Subject: [PATCH 35/36] remove unused import

---
 utils/wandb_logging/wandb_utils.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index 7476000196de..f75b9bcefe6e 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -1,6 +1,5 @@
 import json
 import logging
-import os
 import shutil
 import sys
 from datetime import datetime

From 5651225e2db6ebab33de65047d57870eb88e9147 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 1 Feb 2021 21:28:35 -0800
Subject: [PATCH 36/36] remove unused imports

---
 utils/wandb_logging/wandb_utils.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
index f75b9bcefe6e..264cd4840e3c 100644
--- a/utils/wandb_logging/wandb_utils.py
+++ b/utils/wandb_logging/wandb_utils.py
@@ -1,5 +1,4 @@
 import json
-import logging
 import shutil
 import sys
 from datetime import datetime
@@ -9,7 +8,7 @@
 
 sys.path.append(str(Path(__file__).parent.parent.parent))  # add utils/ to path
 from utils.general import colorstr, xywh2xyxy
-from utils.datasets import img2label_paths
+
 try:
     import wandb
 except ImportError: