ObjDet: Add .yaml and download our trained models

objetos/yolov5/dataset.yaml

@@ -0,0 +1,11 @@
+# Example dataset.yaml file
+
+root: dataset/                   # Path to the root folder | *Optional*
+train: dataset/images/train/     # Path to the training images folder
+val: dataset/images/test/        # Path to the validation images folder
+
+# Number of classes
+nc: 14
+
+# Class names
+names: ['pistol', 'knife', 'bicycle', 'car', 'bus', 'motorbike', 'truck', 'machete', 'pocket/stiletto knife', 'axe', 'rifle/sniper', 'shotgun', 'machine gun', 'submachine gun']

objetos/yolov5/detect_obj.py

@@ -47,7 +47,7 @@ def retrieval(img_path, model_path, output_path, save_as, opt=defaultOpt(), outp
     half = device.type != 'cpu'  # half precision only supported on CUDA
 
     # Load model
-    model = attempt_load(model_path, map_location=device)  # load FP32 model
+    model = attempt_load(model_path, map_location=device, coco_only=opt.coco_only)  # load FP32 model
     stride = int(model.stride.max())  # model stride
     imgsz = check_img_size(imgsz, s=stride)  # check img_size
     names = model.module.names if hasattr(model, 'module') else model.names  # get class names
@@ -166,7 +166,9 @@ def retrieval(img_path, model_path, output_path, save_as, opt=defaultOpt(), outp
     parser.add_argument('--img_path', type=str, default='data/images',
                         help='Source of images/videos. Can be a folder, a single image/video file, '
                              'a url to a image/video, or a json with a list of files (local or remote)')  # file/folder, 0 for webcam
-    parser.add_argument('--model_path', nargs='+', type=str, default='yolov5s.pt', help='model.pt path')
+    parser.add_argument('--model_path', nargs='+', type=str, default='yolov5l.pt', help='model.pt path')
+    parser.add_argument('--coco_only', default=False, action='store_true', 
+                        help='if the model weigths is not found locally, download the models trained only on coco in the original repository.')
 
     parser.add_argument('--output_path', type=str, default='outputs/',
                         help='Path where the outputs generetated will be saved')

objetos/yolov5/hyp.scratch.yaml

@@ -0,0 +1,33 @@
+# Hyperparameters for COCO training from scratch
+# python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300
+# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
+
+
+lr0: 0.01  # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.2  # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937  # SGD momentum/Adam beta1
+weight_decay: 0.0005  # optimizer weight decay 5e-4
+warmup_epochs: 3.0  # warmup epochs (fractions ok)
+warmup_momentum: 0.8  # warmup initial momentum
+warmup_bias_lr: 0.1  # warmup initial bias lr
+box: 0.05  # box loss gain
+cls: 0.5  # cls loss gain
+cls_pw: 1.0  # cls BCELoss positive_weight
+obj: 1.0  # obj loss gain (scale with pixels)
+obj_pw: 1.0  # obj BCELoss positive_weight
+iou_t: 0.20  # IoU training threshold
+anchor_t: 4.0  # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+fl_gamma: 0.0  # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0.015  # image HSV-Hue augmentation (fraction)
+hsv_s: 0.7  # image HSV-Saturation augmentation (fraction)
+hsv_v: 0.4  # image HSV-Value augmentation (fraction)
+degrees: 0.0  # image rotation (+/- deg)
+translate: 0.1  # image translation (+/- fraction)
+scale: 0.5  # image scale (+/- gain)
+shear: 0.0  # image shear (+/- deg)
+perspective: 0.0  # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0  # image flip up-down (probability)
+fliplr: 0.5  # image flip left-right (probability)
+mosaic: 1.0  # image mosaic (probability)
+mixup: 0.0  # image mixup (probability)

objetos/yolov5/models/experimental.py

@@ -110,13 +110,13 @@ def forward(self, x, augment=False):
         return y, None  # inference, train output
 
 
-def attempt_load(weights, map_location=None, inplace=True):
+def attempt_load(weights, map_location=None, inplace=True, coco_only=False):
     from models.yolo import Detect, Model
 
     # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
     model = Ensemble()
     for w in weights if isinstance(weights, list) else [weights]:
-        ckpt = torch.load(attempt_download(w), map_location=map_location)  # load
+        ckpt = torch.load(attempt_download(w, coco_only=coco_only), map_location=map_location)  # load
         model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model
 
     # Compatibility updates

objetos/yolov5/train_light.py

@@ -141,7 +141,7 @@ def train(hyp_path='hyp.scratch.yaml', data='dataset.yaml', output_path='runs/tr
     pretrained = weights.endswith('.pt')
     if pretrained:
         with torch_distributed_zero_first(rank):
-            weights = attempt_download(weights)  # download if not found locally
+            weights = attempt_download(weights, coco_only=opt.coco_only)  # download if not found locally
         ckpt = torch.load(weights, map_location=device)  # load checkpoint
         model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
         exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else []  # exclude keys
@@ -488,6 +488,7 @@ def train(hyp_path='hyp.scratch.yaml', data='dataset.yaml', output_path='runs/tr
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
+    parser.add_argument('--coco_only', default=False, action='store_true', help='if the weights is not found locally, download the models trained only on coco in the original repository.')
     parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
     parser.add_argument('--data', type=str, default='dataset.yaml', help='data.yaml path')
     parser.add_argument('--hyp', type=str, default='hyp.scratch.yaml', help='hyperparameters path')

objetos/yolov5/utils/google_utils.py

@@ -34,10 +34,17 @@ def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
         print('')
 
 
-def attempt_download(file, repo='ultralytics/yolov5'):
+def attempt_download(file, repo='ultralytics/yolov5', coco_only=False):
     # Attempt file download if does not exist
     file = Path(str(file).strip().replace("'", ''))
 
+    if coco_only:
+        repo = 'ultralytics/yolov5'
+        version = 'v5.0'
+    else:
+        repo = 'MPMG-DCC-UFMG/E04'
+        version = 'v1.0'
+
     if not file.exists():
         # URL specified
         name = file.name
@@ -49,23 +56,24 @@ def attempt_download(file, repo='ultralytics/yolov5'):
         # GitHub assets
         file.parent.mkdir(parents=True, exist_ok=True)  # make parent dir (if required)
         try:
-            response = requests.get(f'https://api.github.com/repos/{repo}/releases/download/v5.0').json()  # github api
+            response = requests.get(f'https://api.github.com/repos/{repo}/releases/download/{version}').json()  # github api
             assets = [x['name'] for x in response['assets']]  # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
             tag = response['tag_name']  # i.e. 'v1.0'
         except:  # fallback plan
             assets = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt',
-                      'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
+                    'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
             try:
                 tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
             except:
-                tag = 'v5.0'  # current release
+                tag = 'v5.0' if coco_only else 'v1.0' # current release
 
         if name in assets:
             safe_download(file,
-                          url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
-                          # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}',  # backup url (optional)
-                          min_bytes=1E5,
-                          error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
+                        url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
+                        # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}',  # backup url (optional)
+                        min_bytes=1E5,
+                        error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
+
 
     return str(file)
 

objetos/yolov5/utils/options.py

@@ -24,7 +24,8 @@ def __repr__(self):
 
 class defaultOptTrain():
     def __init__(self):
-        self.weights = 'yolov5s.pt'
+        self.weights = 'yolov5l.pt'
+        self.coco_only = False
         self.cfg = ''
         self.data = 'dataset.yaml'
         self.hyp = 'hyp.scratch.yaml'