def test_read_interlaced_png():
imgs = list(get_images(INTERLACED_PNG, ".png"))
with Image.open(imgs[0]) as im1, Image.open(imgs[1]) as im2:
assert not (im1.info.get("interlace") is im2.info.get("interlace"))
img1 = read_image(imgs[0])
img2 = read_image(imgs[1])
assert_equal(img1, img2)
def test_read_1_bit_png_consistency(shape, mode, tmpdir):
np_rng = np.random.RandomState(0)
image_path = os.path.join(tmpdir, f"test_{shape}.png")
pixels = np_rng.rand(*shape) > 0.5
img = Image.fromarray(pixels)
img.save(image_path)
img1 = read_image(image_path, mode)
img2 = read_image(image_path, mode)
assert_equal(img1, img2)
def test_read_1_bit_png_consistency(shape, mode):
with get_tmp_dir() as root:
image_path = os.path.join(root, f'test_{shape}.png')
pixels = np.random.rand(*shape) > 0.5
img = Image.fromarray(pixels)
img.save(image_path)
img1 = read_image(image_path, mode)
img2 = read_image(image_path, mode)
assert_equal(img1, img2)
def test_read_1_bit_png(shape):
with get_tmp_dir() as root:
image_path = os.path.join(root, f'test_{shape}.png')
pixels = np.random.rand(*shape) > 0.5
img = Image.fromarray(pixels)
img.save(image_path)
img1 = read_image(image_path)
img2 = normalize_dimensions(torch.as_tensor(pixels * 255, dtype=torch.uint8))
assert_equal(img1, img2)
def test_read_1_bit_png(shape, tmpdir):
np_rng = np.random.RandomState(0)
image_path = os.path.join(tmpdir, f"test_{shape}.png")
pixels = np_rng.rand(*shape) > 0.5
img = Image.fromarray(pixels)
img.save(image_path)
img1 = read_image(image_path)
img2 = normalize_dimensions(
torch.as_tensor(pixels * 255, dtype=torch.uint8))
assert_equal(img1, img2)
def test_encode_jpeg(img_path):
img = read_image(img_path)
pil_img = F.to_pil_image(img)
buf = io.BytesIO()
pil_img.save(buf, format="JPEG", quality=75)
encoded_jpeg_pil = torch.frombuffer(buf.getvalue(), dtype=torch.uint8)
for src_img in [img, img.contiguous()]:
encoded_jpeg_torch = encode_jpeg(src_img, quality=75)
assert_equal(encoded_jpeg_torch, encoded_jpeg_pil)
def test_encode_jpeg(img_path):
img = read_image(img_path)
pil_img = F.to_pil_image(img)
buf = io.BytesIO()
pil_img.save(buf, format='JPEG', quality=75)
# pytorch can't read from raw bytes so we go through numpy
pil_bytes = np.frombuffer(buf.getvalue(), dtype=np.uint8)
encoded_jpeg_pil = torch.as_tensor(pil_bytes)
for src_img in [img, img.contiguous()]:
encoded_jpeg_torch = encode_jpeg(src_img, quality=75)
assert_equal(encoded_jpeg_torch, encoded_jpeg_pil)
def visualize_cam_on_img(img_name, model):
cam_extractor = SmoothGradCAMpp(model)
img = read_image(str(img_name))
input_tensor = normalize(
resize(img, (224, 224)) / 255., [0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]).cuda()
out = model(input_tensor.unsqueeze(0))
activation_map = cam_extractor(out.squeeze(0).argmax().item(), out)
result = overlay_mask(to_pil_image(img),
to_pil_image(activation_map, mode='F'),
alpha=0.5)
plt.imshow(result)
plt.axis('off')
plt.tight_layout()
plt.show()
def test_write_jpeg(img_path, tmpdir):
tmpdir = Path(tmpdir)
img = read_image(img_path)
pil_img = F.to_pil_image(img)
torch_jpeg = str(tmpdir / "torch.jpg")
pil_jpeg = str(tmpdir / "pil.jpg")
write_jpeg(img, torch_jpeg, quality=75)
pil_img.save(pil_jpeg, quality=75)
with open(torch_jpeg, "rb") as f:
torch_bytes = f.read()
with open(pil_jpeg, "rb") as f:
pil_bytes = f.read()
assert_equal(torch_bytes, pil_bytes)
def test_write_jpeg(img_path):
with get_tmp_dir() as d:
d = Path(d)
img = read_image(img_path)
pil_img = F.to_pil_image(img)
torch_jpeg = str(d / 'torch.jpg')
pil_jpeg = str(d / 'pil.jpg')
write_jpeg(img, torch_jpeg, quality=75)
pil_img.save(pil_jpeg, quality=75)
with open(torch_jpeg, 'rb') as f:
torch_bytes = f.read()
with open(pil_jpeg, 'rb') as f:
pil_bytes = f.read()
assert_equal(torch_bytes, pil_bytes)
def test_draw_boxes(self):
img = torch.full((3, 100, 100), 255, dtype=torch.uint8)
boxes = torch.tensor(
[[0, 0, 20, 20], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]],
dtype=torch.float)
labels = ["a", "b", "c", "d"]
colors = ["green", "#FF00FF", (0, 255, 0), "red"]
result = utils.draw_bounding_boxes(img,
boxes,
labels=labels,
colors=colors)
path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"assets", "fakedata", "draw_boxes_util.png")
if not os.path.exists(path):
write_png(result, path)
expected = read_image(path)
self.assertTrue(torch.equal(result, expected))
for inputs, labels in validation_data:
inputs = inputs.to(device)
real_val = labels.to(device)
fakes_val = net(inputs)
criterion_loss += criterion(real_val, fakes_val).item()
psnr_score += psnr(real_val, fakes_val).item()
criterion_loss /= validation_size
psnr_score /= validation_size
validation_loss = criterion_loss
val_losses.append(validation_loss)
writer.add_scalar("loss/valid", validation_loss, i)
writer.add_scalar("psnr/valid", psnr_score, i)
speed_mini = read_image(
"speed-mini.png",
mode=ImageReadMode.RGB).to(device).float() / 255.0
writer.add_image("validation image",
net(speed_mini.unsqueeze(0)).squeeze(), i)
print("Validation loss:", validation_loss, "Mean PSNR:",
psnr_score) #, "lr:", scheduler.get_last_lr())
net.train()
if validation_loss < best_loss:
saveNet(filename + "_best", net, optimizer, iterations,
train_losses, val_losses)
print(
f"New best loss: {best_loss} -> {validation_loss}")
best_loss = validation_loss
print("Saved model!")
# This code makes sure that we break both loops if the inner loop is broken out of:
def track(img_dir, model, frame_num, pos_x, pos_y, width, height):
netx = model.netx
netz = model.netz
corr = model.corr
netx.eval()
netz.eval()
corr.eval()
center_x, center_y = pos_x + width / 2, pos_y + height / 2
bboxes = np.zeros((frame_num, 4))
bboxes[0, :] = pos_x, pos_y, width, height
scale_factor = params.scale_step**np.linspace(-(params.scale_num // 2),
(params.scale_num // 2),
params.scale_num)
context = (width + height) * params.context
z_sz = np.sqrt((width + context) * (height + context))
x_sz = float(params.search_sz) / params.examplar_sz * z_sz
min_z = params.scale_min * z_sz
max_z = params.scale_max * z_sz
min_x = params.scale_min * x_sz
max_z = params.scale_max * x_sz
hanning_1d = np.expand_dims(np.hanning(params.final_score_sz), axis=0)
penalty = np.transpose(hanning_1d) * hanning_1d
penalty = penalty / np.sum(penalty)
img_path = os.path.join(img_dir, f'{1:08d}.png')
img_z = read_image(img_path, ImageReadMode.RGB).to(dtype=torch.float32)
frame_sz = img_z.shape
avg_chan = torch.mean(img_z, dim=(-1, -2))
padding_z, x_offset, y_offset = padding_frame(img_z, img_z.shape, center_x,
center_y, z_sz,
avg_chan[0].item())
crop_z = crop_frame_z(img_z, center_x, center_y, x_offset, y_offset, z_sz,
params.examplar_sz)
crop_z = torch.unsqueeze(crop_z, dim=0)
#crop_z = crop_z.to(device)
template_z = netz(crop_z).squeeze(dim=0)
templates_z = torch.stack([template_z] * params.scale_num)
t_start = time.time()
#show_frame(img_z, bboxes[0,:], 1, save_path=os.path.join(img_dir, 'siamFC\\res', f'{1:08d}.png'))
for i in range(1, frame_num):
z_scale = z_sz * scale_factor
x_scale = x_sz * scale_factor
w_scale = width * scale_factor
h_scale = height * scale_factor
img_path = os.path.join(img_dir, f'{i:08d}.png')
img_x = read_image(img_path, ImageReadMode.RGB).to(dtype=torch.float32)
avg_chan = torch.mean(img_z, dim=(-1, -2))
padding_x, x_offset, y_offset = padding_frame(img_x, img_x.shape,
center_x, center_y,
x_scale[-1],
avg_chan[0].item())
x_crops = torch.stack(crop_frame_x(padding_x, center_x, center_y,
x_offset, y_offset, x_scale,
params.search_sz),
dim=0)
#x_crops = x_crops.to(device)
templates_x = netx(x_crops)
scores = corr(templates_x, templates_z)
scores = torchvision.transforms.functional.resize(
scores, (params.final_score_sz, params.final_score_sz),
torchvision.transforms.InterpolationMode.BICUBIC)
scores = torch.squeeze(scores)
for j in range(params.scale_num):
if scale_factor[j] != 1:
scores[j, :, :] = params.scale_penalty * scores[j, :, :]
scale_index = torch.argmax(torch.amax(scores, dim=(-1, -2)))
x_sz = (1 - params.scale_lr
) * x_sz + params.scale_lr * x_scale[scale_index]
width = (1 - params.scale_lr
) * width + params.scale_lr * w_scale[scale_index]
height = (1 - params.scale_lr
) * height + params.scale_lr * h_scale[scale_index]
score = scores[scale_index, :, :]
# print(score.numpy())
score = score - torch.min(score)
score = score / torch.sum(score)
score = (1 - params.window_influence
) * score + params.window_influence * penalty
center_x, center_y = update_target_position(center_x, center_y, score,
x_sz)
center_x = max(0, center_x)
center_x = min(center_x, frame_sz[-1])
center_y = max(0, center_y)
center_y = min(center_y, frame_sz[-2])
bboxes[
i -
1, :] = center_x - width / 2, center_y - height / 2, width, height
#z_sz = (1 - params.scale_lr)*z_sz + params.scale_lr * z_scale[scale_index]
if params.z_lr > 0:
img_z, x_offset, y_offset = padding_frame(img_x, img_x.shape,
center_x, center_y, z_sz,
avg_chan[0].item())
crop_z = crop_frame_z(img_z, center_x, center_y, x_offset,
y_offset, z_sz, params.examplar_sz)
crop_z = torch.unsqueeze(crop_z, dim=0)
new_template_z = netz(crop_z).squeeze()
new_templates_z = torch.stack([new_template_z] * params.scale_num)
templates_z = (
1 - params.z_lr) * templates_z + params.z_lr * new_templates_z
# print(center_x, center_y, width, height)
z_sz = (1 - params.scale_lr
) * z_sz + params.scale_lr * z_scale[scale_index]
# show_frame(img_x, bboxes[i-1,:], i-1)
print(i)
save_path = os.path.join(img_dir, 'siamFC\\res', f'{i:08d}.png')
show_frame(img_x, bboxes[i - 1, :], 1, save_path=save_path)
speed = frame_num / (time.time() - t_start)
plt.close('all')
return bboxes, speed
