""" | """ | ||||
from pathlib import Path | from pathlib import Path | ||||
from models.common import AutoShape | |||||
from models.experimental import attempt_load | from models.experimental import attempt_load | ||||
from models.yolo import Model | from models.yolo import Model | ||||
from utils.downloads import attempt_download | from utils.downloads import attempt_download | ||||
if len(ckpt['model'].names) == classes: | if len(ckpt['model'].names) == classes: | ||||
model.names = ckpt['model'].names # set class names attribute | model.names = ckpt['model'].names # set class names attribute | ||||
if autoshape: | if autoshape: | ||||
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS | |||||
model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS | |||||
return model.to(device) | return model.to(device) | ||||
except Exception as e: | except Exception as e: |
from utils.general import (LOGGER, check_requirements, check_suffix, colorstr, increment_path, make_divisible, | from utils.general import (LOGGER, check_requirements, check_suffix, colorstr, increment_path, make_divisible, | ||||
non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh) | non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh) | ||||
from utils.plots import Annotator, colors, save_one_box | from utils.plots import Annotator, colors, save_one_box | ||||
from utils.torch_utils import time_sync | |||||
from utils.torch_utils import copy_attr, time_sync | |||||
def autopad(k, p=None): # kernel, padding | def autopad(k, p=None): # kernel, padding | ||||
def __init__(self, model): | def __init__(self, model): | ||||
super().__init__() | super().__init__() | ||||
LOGGER.info('Adding AutoShape... ') | |||||
copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes | |||||
self.model = model.eval() | self.model = model.eval() | ||||
def autoshape(self): | |||||
LOGGER.info('AutoShape already enabled, skipping... ') # model already converted to model.autoshape() | |||||
return self | |||||
def _apply(self, fn): | def _apply(self, fn): | ||||
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers | # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers | ||||
self = super()._apply(fn) | self = super()._apply(fn) |
from utils.autoanchor import check_anchor_order | from utils.autoanchor import check_anchor_order | ||||
from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args | from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args | ||||
from utils.plots import feature_visualization | from utils.plots import feature_visualization | ||||
from utils.torch_utils import (copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device, | |||||
time_sync) | |||||
from utils.torch_utils import fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device, time_sync | |||||
try: | try: | ||||
import thop # for FLOPs computation | import thop # for FLOPs computation | ||||
self.info() | self.info() | ||||
return self | return self | ||||
def autoshape(self): # add AutoShape module | |||||
LOGGER.info('Adding AutoShape... ') | |||||
m = AutoShape(self) # wrap model | |||||
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes | |||||
return m | |||||
def info(self, verbose=False, img_size=640): # print model information | def info(self, verbose=False, img_size=640): # print model information | ||||
model_info(self, verbose, img_size) | model_info(self, verbose, img_size) | ||||