def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): """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 Returns: YOLOv5 pytorch model """ set_logging(verbose=verbose) fname = Path(name).with_suffix('.pt') # checkpoint filename try: if pretrained and channels == 3 and classes == 80: model = attempt_load( fname, map_location=torch.device('cpu')) # download/load FP32 model else: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) # create model if pretrained: attempt_download(fname) # download if not found locally ckpt = torch.load(fname, map_location=torch.device('cpu')) # 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 device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available 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 _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): """Creates a specified YOLOv3 model Arguments: name (str): name of model, i.e. 'yolov3' pretrained (bool): load pretrained weights into the model channels (int): number of input channels classes (int): number of model classes autoshape (bool): apply YOLOv3 .autoshape() wrapper to model verbose (bool): print all information to screen device (str, torch.device, None): device to use for model parameters Returns: YOLOv3 pytorch model """ from pathlib import Path from models.yolo import Model, attempt_load from utils.general import check_requirements, set_logging from utils.google_utils import attempt_download from utils.torch_utils import select_device check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop')) set_logging(verbose=verbose) fname = Path(name).with_suffix('.pt') # checkpoint filename try: if pretrained and channels == 3 and classes == 80: model = attempt_load(fname, map_location=torch.device('cpu')) # download/load FP32 model else: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) # create model if pretrained: ckpt = torch.load(attempt_download(fname), map_location=torch.device('cpu')) # 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 device = select_device('0' if torch.cuda.is_available() else 'cpu') if device is None else torch.device(device) 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
# ------------------------- yolov5s 모델 import ------------------------- from pathlib import Path from models.yolo import Model, attempt_load from utils.general import check_requirements, set_logging from utils.google_utils import attempt_download from utils.torch_utils import select_device check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop')) # set_logging(verbose=verbose) device = None name1 = 'yolov5_custom_pink.pt' fname1 = Path(name1).with_suffix('.pt') # checkpoint filename model1 = attempt_load(fname1, map_location=torch.device('cpu')) model1 = model1.autoshape() # for file/URI/PIL/cv2/np inputs and NMS model1.conf = 0.5 name2 = 'yolov5_custom.pt' fname2 = Path(name2).with_suffix('.pt') # checkpoint filename model2 = attempt_load(fname2, map_location=torch.device('cpu')) model2 = model2.autoshape() # for file/URI/PIL/cv2/np inputs and NMS model2.conf = 0.5 device = select_device('0' if torch.cuda.is_available() else 'cpu' ) if device is None else torch.device(device) # ------------------------- 모델 import fin ------------------------- """ on_connect는 subscriber가 브로커에 연결하면서 호출할 함수 rc가 0이면 정상접속이 됐다는 의미