Hi erik !
I've been working on this topic and came with the code below to generate the .blob file to perform the concatenation of two images. Do you have any suggestions on how to proceed to process (concatenate with the code below) two images let's say 30 frames apart in time? Should I use some kind of queue to process them?
Thanks once again!
from pathlib import Path
import torch
from torch import nn
import kornia
import onnx
from onnxsim import simplify
import blobconverter
name = 'concat'
class Model(nn.Module):
def forward(self, img1, img2):
# Transforming to B&W
img1bw = kornia.color.rgb_to_grayscale(img1, rgb_weights=torch.tensor([0.299, 0.587, 0.114]))
img2bw = kornia.color.rgb_to_grayscale(img2, rgb_weights=torch.tensor([0.299, 0.587, 0.114]))
# Global Parameters - Screening
band = 300 # Band of height to be analyzed for the images (# of pixels)
int_ref = 20 # Pixel range +/- to be evaluated for the fine analysis
end = img1bw.shape[2] - band
val = 0
reduction = int_ref # Resolution reduction
val = self.screen(band, val, end, reduction, img2bw, img1bw)
reduction = 1 # Resolution reduction
val = self.screen(band, val - int_ref, val + int_ref, reduction, img2bw, img1bw)
img1cut = torch.narrow(img1, 2, 0, val)
imgfin = torch.cat((img1cut, img2), 2)
return imgfin
def screen(self, band, start, end, imrdc, imgref, imgscr):
rg = range(start, end, 1) # Range for y0 - starting point for image being screened
diff = [None] * len([*rg])
i = 0
# Band of reference
ref_gray = torch.narrow(imgref, 2, 0, band)
ref_gray = kornia.geometry.transform.rescale(ref_gray, (1/imrdc, 1/imrdc))
for y0 in rg:
# Band of comparison
gray = torch.narrow(imgscr, 2, y0, band)
gray = kornia.geometry.transform.rescale(gray, (1/imrdc, 1 /imrdc))
# Absolute difference between bands
change = torch.abs(torch.sub(gray, ref_gray))
diff[i] = torch.sum(change) / (change.shape[2] * change.shape[3]) # Average difference measure
i = i + 1
val = [*rg][diff.index(min(diff))] # Cut height
if val == 0:
val = 100 # .blob is not generated if val == 0
return val
# Define the expected input shape (dummy input)
shape = (1, 3, 1080, 1920)
model = Model()
X = torch.ones(shape, dtype=torch.float32)
path = Path("out/")
path.mkdir(parents=True, exist_ok=True)
onnx_path = str(path / (name + '.onnx'))
print(f"Writing to {onnx_path}")
torch.onnx.export(
model,
(X, X),
onnx_path,
opset_version=12,
do_constant_folding=True,
)
onnx_simplified_path = str(path / (name + '_simplified.onnx'))
# Use onnx-simplifier to simplify the onnx model
onnx_model = onnx.load(onnx_path)
model_simp, check = simplify(onnx_model)
onnx.save(model_simp, onnx_simplified_path)
# Use blobconverter to convert onnx->IR->blob
blobconverter.from_onnx(
model=onnx_simplified_path,
data_type="FP16",
shaves=6,
use_cache=False,
output_dir=path,
optimizer_params=[]
)
