| 243 | |
| 244 | @threaded |
| 245 | def plot_images(images, targets, paths=None, fname='images.jpg', names=None): |
| 246 | # Plot image grid with labels |
| 247 | if isinstance(images, torch.Tensor): |
| 248 | images = images.cpu().float().numpy() |
| 249 | if isinstance(targets, torch.Tensor): |
| 250 | targets = targets.cpu().numpy() |
| 251 | |
| 252 | max_size = 1920 # max image size |
| 253 | max_subplots = 16 # max image subplots, i.e. 4x4 |
| 254 | bs, _, h, w = images.shape # batch size, _, height, width |
| 255 | bs = min(bs, max_subplots) # limit plot images |
| 256 | ns = np.ceil(bs ** 0.5) # number of subplots (square) |
| 257 | if np.max(images[0]) <= 1: |
| 258 | images *= 255 # de-normalise (optional) |
| 259 | |
| 260 | # Build Image |
| 261 | mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init |
| 262 | for i, im in enumerate(images): |
| 263 | if i == max_subplots: # if last batch has fewer images than we expect |
| 264 | break |
| 265 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin |
| 266 | im = im.transpose(1, 2, 0) |
| 267 | mosaic[y:y + h, x:x + w, :] = im |
| 268 | |
| 269 | # Resize (optional) |
| 270 | scale = max_size / ns / max(h, w) |
| 271 | if scale < 1: |
| 272 | h = math.ceil(scale * h) |
| 273 | w = math.ceil(scale * w) |
| 274 | mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h))) |
| 275 | |
| 276 | # Annotate |
| 277 | fs = int((h + w) * ns * 0.01) # font size |
| 278 | annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names) |
| 279 | for i in range(i + 1): |
| 280 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin |
| 281 | annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders |
| 282 | if paths: |
| 283 | annotator.text((x + 5, y + 5), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames |
| 284 | if len(targets) > 0: |
| 285 | ti = targets[targets[:, 0] == i] # image targets |
| 286 | boxes = xywh2xyxy(ti[:, 2:6]).T |
| 287 | classes = ti[:, 1].astype('int') |
| 288 | labels = ti.shape[1] == 6 # labels if no conf column |
| 289 | conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred) |
| 290 | |
| 291 | if boxes.shape[1]: |
| 292 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01 |
| 293 | boxes[[0, 2]] *= w # scale to pixels |
| 294 | boxes[[1, 3]] *= h |
| 295 | elif scale < 1: # absolute coords need scale if image scales |
| 296 | boxes *= scale |
| 297 | boxes[[0, 2]] += x |
| 298 | boxes[[1, 3]] += y |
| 299 | for j, box in enumerate(boxes.T.tolist()): |
| 300 | cls = classes[j] |
| 301 | color = colors(cls) |
| 302 | cls = names[cls] if names else cls |