def done(self):
Path(self.directory).mkdir(parents=True, exist_ok=True)
stream = torch.from_numpy(np.stack(self.t))
for image in stream:
write_jpeg(image.permute(2, 0, 1),
f'{self.directory}/{self.image_id}.jpg')
self.image_id += 1
def overlay_img_with_density(img_path, density_map_path, output_path):
"""
combine output density map with image to create the red heatmap overlay
:param img_path:
:param density_map_path: output .torch of density map
:param output_path:
:return:
"""
img_tensor = read_image(img_path)
density_map_tensor = torch.load(density_map_path)
print(img_tensor.shape)
print(density_map_tensor.shape)
print(density_map_tensor.sum())
density_map_tensor = torch.from_numpy(density_map_tensor).unsqueeze(
dim=0).unsqueeze(dim=0)
print("density_map_tensor.shape",
density_map_tensor.shape) # torch.Size([1, 1, 46, 82])
upsampling_density_map_tensor = nn.functional.interpolate(
density_map_tensor, scale_factor=8) / 64
overlay_density_map = img_tensor.detach().clone()
upsampling_density_map_tensor = (
upsampling_density_map_tensor.squeeze(dim=0) /
upsampling_density_map_tensor.max() * 255)
overlay_density_map[0] = torch.clamp_max(
img_tensor[0] + upsampling_density_map_tensor[0] * 2, max=255)
write_jpeg(overlay_density_map.type(torch.uint8), output_path, quality=100)
def imsave_tensor(imgtensor, path):
"""
Input a CHW tensor in float between 0 and 1
Save into *.jpg file
"""
if imgtensor.device.type != 'cpu':
imgtensor = imgtensor.to('cpu')
imgtensor = imgtensor * 255
imgtensor = imgtensor.byte()
write_jpeg(imgtensor, path)
def overlay_img_with_density_padding(img_origin, density_map_path,
output_path):
"""
:param img_path:
:param density_map_path: output .torch of density map
:param output_path:
:return:
"""
# # img_origin = Image.open(img_path).convert('RGB')
#
# transformer = transforms.Compose([
# transforms.ToTensor()
# ])
img_tensor = read_image(img_origin)
print("img tensor shape", img_tensor.shape)
density_map_tensor = torch.load(density_map_path)
print(img_tensor.shape)
print(density_map_tensor.shape)
print(density_map_tensor.sum())
density_map_tensor = torch.from_numpy(density_map_tensor).unsqueeze(
dim=0).unsqueeze(dim=0)
print("density_map_tensor.shape",
density_map_tensor.shape) # torch.Size([1, 1, 46, 82])
upsampling_density_map_tensor = nn.functional.interpolate(
density_map_tensor, scale_factor=8) / 64
pad_density_map_tensor = torch.zeros(
(1, 3, img_tensor.shape[1], img_tensor.shape[2]))
pad_density_map_tensor[:, 0, :upsampling_density_map_tensor.
shape[2], :upsampling_density_map_tensor.
shape[3]] = upsampling_density_map_tensor
overlay_density_map = img_tensor.detach().clone()
pad_density_map_tensor = (pad_density_map_tensor.squeeze(dim=0) /
pad_density_map_tensor.max() * 255)
overlay_density_map[0] = torch.clamp_max(img_tensor[0] +
pad_density_map_tensor[0] * 2,
max=255)
write_jpeg(overlay_density_map.type(torch.uint8), output_path, quality=100)
def track_frames(frames, start, step):
frames = th.stack(frames)
# Track person only
with th.cuda.amp.autocast(enabled=cfg.det_amp):
dets, features = detector.detect(frames,
size=cfg.det_scales,
conf_thres=cfg.det_cls_thres,
iou_thres=cfg.det_nms_thres,
batch_preprocess=True)
persons = dets_select(dets, cfg.trk_cls_person)
objs = [
dets_f[~persons_f].cpu()
for dets_f, persons_f in zip(dets, persons)
]
ppls = [
dets_f[persons_f].cpu()
for dets_f, persons_f in zip(dets, persons)
]
ppl_feats = [
feats_f[persons_f].cpu()
for feats_f, persons_f in zip(features, persons)
]
for j, (objs_f, ppls_f,
ppl_feats_f) in enumerate(zip(objs, ppls, ppl_feats), start):
logging.info(
f"[{start + (j - start) * step}] objs: {tuple(objs_f.shape)}, ppls: {tuple(ppls_f.shape)}, feats: {tuple(ppl_feats_f.shape)}"
)
assert objs_f.shape[1] == 4 + 1 + 1
assert ppls_f.shape[1] == 4 + 1 + 1
assert len(ppls) == len(ppl_feats)
# assert ppl_feats.shape[1] == 256 + 512 + 1024
assert ppl_feats_f.shape[1] == 320 + 640 + 1280
matches = tracker.update(
ppls_f,
ppl_feats_f.view(len(ppl_feats_f),
np.prod(ppl_feats_f.shape[1:])))
snapshot = tracker.snapshot()
tracks = []
for tid, info in snapshot:
tracks.append([tid, info])
logging.debug(f"matches[{start + (j - start) * step}]: {matches}")
logging.debug(
f"snapshot[{start + (j - start) * step}]: {snapshot}")
# Render both dets and tracks side by side
frame_det = frames[j - start]
frame_trk = frame_det.clone()
C, H, W = frame_det.shape
idx = f'{start + (j - start) * step:03d}'
if cfg.render_all:
dets = th.cat([ppls_f, objs_f])
frame_det = render_frame(idx, frame_det, dets, False)
else:
frame_det = render_frame(idx, frame_det, ppls_f, False)
if tracks:
frame_trk = render_frame(idx, frame_trk, tracks, True)
frame = th.zeros((C, H, 2 * W), dtype=th.uint8)
frame[:, :, :W] = frame_det
frame[:, :, W:] = frame_trk
write_jpeg(frame, str(export_frame / f"frame{idx}.jpg"))
if media is not None:
frame = av.VideoFrame.from_ndarray(frame.permute(1, 2,
0).numpy(),
format='rgb24')
packets = stream.encode(frame)
media.mux(packets)
logging.debug(f'Encoded: {len(packets)} {packets}, {frame}')
train = datasets.MNIST("../../data", train=True, download=False)
print(f"raw_folder: test: {test.raw_folder}, train: {train.raw_folder}")
print(
f"processed_folder: test: {test.processed_folder}, train: {train.processed_folder}"
)
print(f"extra_repr:\ntest: {test.extra_repr}\ntrain: {train.extra_repr}")
print(f"class to index: {test.class_to_idx}")
# 读写图像: torchvision.io包
tensor = io.read_image("../../data/image/1.jpg")
print("tensor shape:", tensor.shape)
io.write_png(tensor, "../../data/image/result.png")
tensor = io.read_image("../../data/image/lena.png")
print("tensor shape:", tensor.shape)
io.write_jpeg(tensor, "../../data/image/result.jpg")
# 下载pre-trained AlexNet模型: torchvision.models包
net = models.alexnet(pretrained=True)
# 计算机视觉操作: torchvision.ops包
boxes = torch.tensor([[1, 1, 101, 101], [3, 5, 13, 15], [2, 4, 22, 44]])
area = ops.box_area(boxes)
print(f"area: {area}")
index = ops.remove_small_boxes(boxes, min_size=20)
print(f"index: {index}")
# 图像变换: torchvision.transforms包
resize = transforms.Resize(size=[256, 128])
img = resize.forward(tensor)