diff --git a/__init__.py b/__init__.py index 57fde11..335d3b4 100644 --- a/__init__.py +++ b/__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() diff --git a/models/experimental.py b/models/experimental.py index d1dd34c..717fe4c 100644 --- a/models/experimental.py +++ b/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 diff --git a/utils/general.py b/utils/general.py index 298db1d..4de6c48 100644 --- a/utils/general.py +++ b/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 diff --git a/utils/google_utils.py b/utils/google_utils.py new file mode 100644 index 0000000..0a7ca3b --- /dev/null +++ b/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))