diff --git a/detect.py b/detect.py index 2fb96eaa59a1..9d922aed1d12 100644 --- a/detect.py +++ b/detect.py @@ -46,7 +46,7 @@ def detect(save_img=False): dataset = LoadImages(source, img_size=imgsz) # Get names and colors - names = model.names if hasattr(model, 'names') else model.modules.names + names = model.module.names if hasattr(model, 'module') else model.names colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] # Run inference diff --git a/train.py b/train.py index d9d2884d6e97..ea72f1549f9a 100644 --- a/train.py +++ b/train.py @@ -79,6 +79,7 @@ def train(hyp): # Create model model = Model(opt.cfg).to(device) assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc']) + model.names = data_dict['names'] # Image sizes gs = int(max(model.stride)) # grid size (max stride) @@ -193,7 +194,6 @@ def train(hyp): model.hyp = hyp # attach hyperparameters to model model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou) model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights - model.names = data_dict['names'] # Class frequency labels = np.concatenate(dataset.labels, 0)