Spaces:
Sleeping
Sleeping
File size: 1,860 Bytes
81ecb2b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import cv2
def label_image(
image,
label,
font_scale=1.0,
font_thickness=1,
label_origin=(10, 64),
font_color=(255, 255, 255),
font=cv2.FONT_HERSHEY_SIMPLEX,
):
text_size, baseline = cv2.getTextSize(label, font, font_scale, font_thickness)
image[
label_origin[1] - text_size[1] : label_origin[1] + baseline,
label_origin[0] : label_origin[0] + text_size[0],
] = (255 - font_color[0], 255 - font_color[1], 255 - font_color[2])
cv2.putText(
image, label, label_origin, font, font_scale, font_color, font_thickness
)
return image
def to_device(values, device=None, non_blocking=True):
"""Transfer a set of values to the device.
Args:
values: a nested dict/list/tuple of tensors
device: argument to `to()` for the underlying vector
NOTE:
if the device is not specified, using `th.cuda()`
"""
if device is None:
device = th.device("cuda")
if isinstance(values, dict):
return {k: to_device(v, device=device) for k, v in values.items()}
elif isinstance(values, tuple):
return tuple(to_device(v, device=device) for v in values)
elif isinstance(values, list):
return [to_device(v, device=device) for v in values]
elif isinstance(values, th.Tensor):
return values.to(device, non_blocking=non_blocking)
elif isinstance(values, nn.Module):
return values.to(device)
elif isinstance(values, np.ndarray):
return th.from_numpy(values).to(device)
else:
return values
|