# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
PyTorch Hub yolo_models https://pytorch.org/hub/ultralytics_yolov5/
Usage:
	import torch
	model = torch.hub.load('ultralytics/models', 'yolov5s')
"""

import torch

def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	"""Creates a specified YOLOv5 model
	Arguments:
		name (str): name of model, i.e. 'yolov5s'
		pretrained (bool): load pretrained weights into the model
		channels (int): number of input channels
		classes (int): number of model classes
		autoshape (bool): apply YOLOv5 .autoshape() wrapper to model
		verbose (bool): print all information to screen
		device (str, torch.device, None): device to use for model parameters
	Returns:
		YOLOv5 pytorch model
	"""
	from pathlib import Path
	import os
	from soco_yolo.yolo import Model
	from soco_yolo.experimental import attempt_load
	from soco_yolo.utils.general import check_requirements, set_logging
	from soco_yolo.utils.downloads import attempt_download
	from soco_yolo.utils.torch_utils import select_device
	from soco_yolo.config import EnvVar

	check_requirements(exclude=('tensorboard', 'thop', 'opencv-python'))
	set_logging(verbose=verbose)

	# save_dir = Path('') if str(name).endswith('.pt') else file.parent
	root_dir = os.path.dirname(os.path.abspath(__file__)).replace('soco_yolo', EnvVar.MODEL_DIR)
	save_dir = Path(root_dir).resolve()
	path = (save_dir / name).with_suffix('.pt')  # checkpoint path
	try:
		device = EnvVar.DEVICE
		if pretrained and channels == 3 and classes == 80:
			#model = attempt_load(path, map_location=device)  # download/load FP32 model
			model = attempt_load(name, map_location=device)  # download/load FP32 model
		else:
			cfg = list((Path(__file__).parent / 'yolo_models').rglob(f'{name}.yaml'))[0]  # model.yaml path
			model = Model(cfg, channels, classes)  # create model
			if pretrained:
				ckpt = torch.load(attempt_download(path), map_location=device)	# load
				msd = model.state_dict()  # model state_dict
				csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
				csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape}  # filter
				model.load_state_dict(csd, strict=False)  # load
				if len(ckpt['model'].names) == classes:
					model.names = ckpt['model'].names  # set class names attribute
		if autoshape:
			model = model.autoshape()  # for file/URI/PIL/cv2/np inputs and NMS

		return model.to(device)

	except Exception as e:
		help_url = 'https://github.com/ultralytics/yolov5/issues/36'
		s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url
		raise Exception(s) from e


def custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None):
	# YOLOv5 custom or local model
	return _create(path, autoshape=autoshape, verbose=verbose, device=device)


def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-small model https://github.com/ultralytics/yolov5
	return _create('yolov5s', pretrained, channels, classes, autoshape, verbose, device)


def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-medium model https://github.com/ultralytics/yolov5
	return _create('yolov5m', pretrained, channels, classes, autoshape, verbose, device)


def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-large model https://github.com/ultralytics/yolov5
	return _create('yolov5l', pretrained, channels, classes, autoshape, verbose, device)


def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-xlarge model https://github.com/ultralytics/yolov5
	return _create('yolov5x', pretrained, channels, classes, autoshape, verbose, device)


def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
	return _create('yolov5s6', pretrained, channels, classes, autoshape, verbose, device)


def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
	return _create('yolov5m6', pretrained, channels, classes, autoshape, verbose, device)


def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
	return _create('yolov5l6', pretrained, channels, classes, autoshape, verbose, device)


def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
	# YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
	return _create('yolov5x6', pretrained, channels, classes, autoshape, verbose, device)


if __name__ == '__main__':
	model = _create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True)	# pretrained
	# model = custom(path='path/to/model.pt')  # custom

	# Verify inference
	import cv2
	import numpy as np
	from PIL import Image
	from pathlib import Path

	imgs = ['data/images/zidane.jpg',  # filename
			Path('data/images/zidane.jpg'),  # Path
			'https://ultralytics.com/images/zidane.jpg',  # URI
			cv2.imread('data/images/bus.jpg')[:, :, ::-1],	# OpenCV
			Image.open('data/images/bus.jpg'),	# PIL
			np.zeros((320, 640, 3))]  # numpy

	results = model(imgs)  # batched inference
	results.print()
	results.save()
