diff --git a/utils/general.py b/utils/general.py index 8bbf76970966..11a2cf42f9f7 100755 --- a/utils/general.py +++ b/utils/general.py @@ -1,5 +1,6 @@ import glob import logging +import math import os import platform import random @@ -12,7 +13,6 @@ from pathlib import Path import cv2 -import math import matplotlib import matplotlib.pyplot as plt import numpy as np @@ -589,7 +589,7 @@ def build_targets(p, targets, model): # Append a = t[:, 6].long() # anchor indices - indices.append((b, a, gj, gi)) # image, anchor, grid indices + indices.append((b, a, gj.clamp_(0, gain[3]), gi.clamp_(0, gain[2]))) # image, anchor, grid indices tbox.append(torch.cat((gxy - gij, gwh), 1)) # box anch.append(anchors[a]) # anchors tcls.append(c) # class @@ -1096,7 +1096,8 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3) if fname is not None: - mosaic = cv2.resize(mosaic, (int(ns * w * 0.5), int(ns * h * 0.5)), interpolation=cv2.INTER_AREA) + r = min(1280. / max(h, w) / ns, 1.0) # ratio to limit image size + mosaic = cv2.resize(mosaic, (int(ns * w * r), int(ns * h * r)), interpolation=cv2.INTER_AREA) # cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB)) # cv2 save Image.fromarray(mosaic).save(fname) # PIL save return mosaic