fix: predict result

__init__.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 
 from models.yolo import Model
+from models.experimental import attempt_load
 from utils.torch_utils import select_device, intersect_dicts
 from utils.datasets import letterbox
 from utils.general import scale_coords, non_max_suppression
@@ -59,13 +60,14 @@ def define(hyp, opt, device, recoverPath):
   ckpt = torch.load(modelPath, map_location=device)
   if hyp.get('anchors'):
       ckpt['model'].yaml['anchors'] = round(hyp['anchors'])
+  predict_model = attempt_load(modelPath, map_location=device)
   model = Model(os.path.join(os.path.dirname(__file__), 'models', 'yolov5s.yaml'), ch=3, nc=opt.nc).to(device)
   exclude = ['anchor']
   state_dict = ckpt['model'].float().state_dict()  # to FP32
   state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
   model.load_state_dict(state_dict, strict=False)
   logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
-  return model, ckpt
+  return model, predict_model, ckpt
 
 def main(data, args):
   opt = obj({})
@@ -88,10 +90,11 @@ def main(data, args):
   with open(opt.hyp) as f:
       hyp = yaml.load(f, Loader=yaml.FullLoader)
   device = select_device(opt.device, batch_size=opt.batch_size)
-  yolov5, ckpt = define(hyp, opt, device, recoverPath)
+  yolov5, predict_model, ckpt = define(hyp, opt, device, recoverPath)
   half = device.type != 'cpu'
   class PipcookModel:
     model = yolov5
+    p_model = predict_model
     config = {
       "ckpt": ckpt,
       "img_size": opt.img_size
@@ -109,9 +112,26 @@ def predict(self, inputData):
           img = img.unsqueeze(0)
 
       # Inference
-      pred = self.model(img)[0]
+      pred = self.p_model(img)[0]
       pred = non_max_suppression(pred, 0.25, 0.45)
-      return pred
+
+      # Parse Inference
+      boxes = []
+      classes = []
+      scores = []
+      for i, det in enumerate(pred):  # detections per image
+        # Write results
+        for *xyxy, conf, cls in reversed(det):
+          boxes.append([int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])])
+          classes.append(int(cls))
+          scores.append(float(conf))
+      output = {
+        'boxes': boxes,
+        'classes': classes,
+        'scores': scores
+      }
+      return output
+
   sys.path.pop()
   sys.path.pop()
   return PipcookModel()

models/experimental.py

@@ -5,7 +5,7 @@
 import torch.nn as nn
 
 from .common import Conv, DWConv
-
+from utils.google_utils import attempt_download
 
 class CrossConv(nn.Module):
     # Cross Convolution Downsample
@@ -126,3 +126,25 @@ def forward(self, x, augment=False):
         # y = torch.cat(y, 1)  # nms ensemble
         y = torch.stack(y).mean(0)  # mean ensemble
         return y, None  # inference, train output
+
+def attempt_load(weights, map_location=None):
+    # 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]:
+        attempt_download(w)
+        model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval())  # load FP32 model
+
+    # Compatibility updates
+    for m in model.modules():
+        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
+            m.inplace = True  # pytorch 1.7.0 compatibility
+        elif type(m) is Conv:
+            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+
+    if len(model) == 1:
+        return model[-1]  # return model
+    else:
+        print('Ensemble created with %s\n' % weights)
+        for k in ['names', 'stride']:
+            setattr(model, k, getattr(model[-1], k))
+        return model  # return ensemble

utils/general.py

@@ -3,7 +3,8 @@
 import numpy as np
 import math
 import time
-
+import torch
+import torchvision
 from .torch_utils import init_torch_seeds
 
 
@@ -17,6 +18,15 @@ def init_seeds(seed=0):
   np.random.seed(seed)
   init_torch_seeds(seed)
 
+def xywh2xyxy(x):
+    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+    return y
+
 def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False):
     """Performs Non-Maximum Suppression (NMS) on inference results
 

utils/google_utils.py

@@ -0,0 +1,122 @@
+# Google utils: https://cloud.google.com/storage/docs/reference/libraries
+
+import os
+import platform
+import subprocess
+import time
+from pathlib import Path
+
+import requests
+import torch
+
+
+def gsutil_getsize(url=''):
+    # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
+    s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8')
+    return eval(s.split(' ')[0]) if len(s) else 0  # bytes
+
+
+def attempt_download(file, repo='ultralytics/yolov5'):
+    # Attempt file download if does not exist
+    file = Path(str(file).strip().replace("'", '').lower())
+
+    if not file.exists():
+        try:
+            response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').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']
+            tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
+
+        name = file.name
+        if name in assets:
+            msg = f'{file} missing, try downloading from https://github.com/{repo}/releases/'
+            redundant = False  # second download option
+            try:  # GitHub
+                url = f'https://github.com/{repo}/releases/download/{tag}/{name}'
+                print(f'Downloading {url} to {file}...')
+                torch.hub.download_url_to_file(url, file)
+                assert file.exists() and file.stat().st_size > 1E6  # check
+            except Exception as e:  # GCP
+                print(f'Download error: {e}')
+                assert redundant, 'No secondary mirror'
+                url = f'https://storage.googleapis.com/{repo}/ckpt/{name}'
+                print(f'Downloading {url} to {file}...')
+                os.system(f'curl -L {url} -o {file}')  # torch.hub.download_url_to_file(url, weights)
+            finally:
+                if not file.exists() or file.stat().st_size < 1E6:  # check
+                    file.unlink(missing_ok=True)  # remove partial downloads
+                    print(f'ERROR: Download failure: {msg}')
+                print('')
+                return
+
+
+def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
+    # Downloads a file from Google Drive. from yolov5.utils.google_utils import *; gdrive_download()
+    t = time.time()
+    file = Path(file)
+    cookie = Path('cookie')  # gdrive cookie
+    print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='')
+    file.unlink(missing_ok=True)  # remove existing file
+    cookie.unlink(missing_ok=True)  # remove existing cookie
+
+    # Attempt file download
+    out = "NUL" if platform.system() == "Windows" else "/dev/null"
+    os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}')
+    if os.path.exists('cookie'):  # large file
+        s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}'
+    else:  # small file
+        s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"'
+    r = os.system(s)  # execute, capture return
+    cookie.unlink(missing_ok=True)  # remove existing cookie
+
+    # Error check
+    if r != 0:
+        file.unlink(missing_ok=True)  # remove partial
+        print('Download error ')  # raise Exception('Download error')
+        return r
+
+    # Unzip if archive
+    if file.suffix == '.zip':
+        print('unzipping... ', end='')
+        os.system(f'unzip -q {file}')  # unzip
+        file.unlink()  # remove zip to free space
+
+    print(f'Done ({time.time() - t:.1f}s)')
+    return r
+
+
+def get_token(cookie="./cookie"):
+    with open(cookie) as f:
+        for line in f:
+            if "download" in line:
+                return line.split()[-1]
+    return ""
+
+# def upload_blob(bucket_name, source_file_name, destination_blob_name):
+#     # Uploads a file to a bucket
+#     # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
+#
+#     storage_client = storage.Client()
+#     bucket = storage_client.get_bucket(bucket_name)
+#     blob = bucket.blob(destination_blob_name)
+#
+#     blob.upload_from_filename(source_file_name)
+#
+#     print('File {} uploaded to {}.'.format(
+#         source_file_name,
+#         destination_blob_name))
+#
+#
+# def download_blob(bucket_name, source_blob_name, destination_file_name):
+#     # Uploads a blob from a bucket
+#     storage_client = storage.Client()
+#     bucket = storage_client.get_bucket(bucket_name)
+#     blob = bucket.blob(source_blob_name)
+#
+#     blob.download_to_filename(destination_file_name)
+#
+#     print('Blob {} downloaded to {}.'.format(
+#         source_blob_name,
+#         destination_file_name))