[Python] Complete a picture of a complete curve

I have a pair of images in the data/in/0.png and data/out/0.png folders like the two images opposite. I generated them programmatically in Python.



Here is the code to generate the pictures.

import copy
import matplotlib.pyplot as plt
import numpy as np
import random
import scipy
import cv2

def random_spaced_points(n, a = 0.0, b = 1.0, delta = 0.1, rng = None):
if (rng is None):
rng = np.random.default_rng()

if (b - a < (n - 1) * delta):
raise ValueError("Intervalle trop petit pour placer tous les points avec cet espacement.")

base = rng.uniform(a, b - (n - 1) * delta, size=n)
base.sort()

spaced = base + np.arange(n) * delta

return spaced

def drawPolynom(img, p, x0, x1, n = 100):
x = np.linspace(x0, x1, 200)
y = p(x)

coords = list(zip([int(z) for z in x], [int(z) for z in y]))

for i in range(1, len(coords)):
cv2.line(img, coords[i - 1], coords, (255, 0, 0), 1)

return img

for i in range(0, 100):
size = (800, 600)

channels = 3 # 3 pour une image couleur (BGR)
img = np.zeros((size[1], size[0], channels), dtype = np.uint8)

n = 4
alpha = 0.25
x = sorted(list(random_spaced_points(n, alpha * size[0], (1 - alpha) * size[0], alpha / 2 * size[0])))
y = list(random_spaced_points(n, alpha * size[1], (1 - alpha) * size[1], alpha / 2 * size[1]))
random.shuffle(y)

p = np.polynomial.Polynomial.fit(x, y, deg = 3)

img1 = copy.deepcopy(img)
drawPolynom(img1, p, x[0], x[-1])

"""
plt.imshow(img1)
#plt.plot(x, y)
mng = plt.get_current_fig_manager()
mng.window.state('zoomed')
plt.show()
"""
cv2.imwrite("data/out/" + str(i) + ".png", cv2.cvtColor(img1, cv2.COLOR_RGB2BGR))

img2 = copy.deepcopy(img)
alpha = 1 / 3
drawPolynom(img2, p, x[0], (1 - alpha) * x[0] + alpha * x[-1])
drawPolynom(img2, p, alpha * x[0] + (1 - alpha) * x[-1], x[-1])

"""
plt.imshow(img2)
mng = plt.get_current_fig_manager()
mng.window.state('zoomed')
plt.show()
"""
cv2.imwrite("data/in/" + str(i) + ".png", cv2.cvtColor(img2, cv2.COLOR_RGB2BGR))


I want to use a neural network to complete the image data/in/*.png to find the complete curve of data/out/*.png. I'm new to this and asked ChatGPT for help. Here's the code he suggested, but it doesn't work.

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import glob

class CurveDataset(Dataset):
def __init__(self, in_dir, out_dir, transform=None):
self.in_files = sorted(glob.glob(f"{in_dir}/*.png"))
self.out_files = sorted(glob.glob(f"{out_dir}/*.png"))
self.transform = transform

def __len__(self):
return len(self.in_files)

def __getitem__(self, idx):
img_in = Image.open(self.in_files[idx]).convert("L")
img_out = Image.open(self.out_files[idx]).convert("L")

if self.transform:
img_in = self.transform(img_in)
img_out = self.transform(img_out)

return img_in, img_out

transform = transforms.Compose([
transforms.Resize((128,128)), # normalisation de taille
transforms.ToTensor() # convertit en [0,1]
])

import torch.nn as nn

class DoubleConv(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.conv(x)

class UNet(nn.Module):
def __init__(self):
super().__init__()
self.down1 = DoubleConv(1, 32)
self.pool1 = nn.MaxPool2d(2)
self.down2 = DoubleConv(32, 64)
self.pool2 = nn.MaxPool2d(2)
self.down3 = DoubleConv(64, 128)
self.pool3 = nn.MaxPool2d(2)

self.mid = DoubleConv(128, 256)

self.up3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
self.conv3 = DoubleConv(256, 128)
self.up2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
self.conv2 = DoubleConv(128, 64)
self.up1 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
self.conv1 = DoubleConv(64, 32)

self.out_conv = nn.Conv2d(32, 1, kernel_size=1)

def forward(self, x):
d1 = self.down1(x)
d2 = self.down2(self.pool1(d1))
d3 = self.down3(self.pool2(d2))
m = self.mid(self.pool3(d3))
u3 = self.up3(m)
u3 = self.conv3(torch.cat([u3, d3], dim=1))
u2 = self.up2(u3)
u2 = self.conv2(torch.cat([u2, d2], dim=1))
u1 = self.up1(u2)
u1 = self.conv1(torch.cat([u1, d1], dim=1))
return torch.sigmoid(self.out_conv(u1))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# suppose que tu as ./data/in et ./data/out
train_dataset = CurveDataset("data/in", "data/out", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)

model = UNet().to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(5):
for img_in, img_out in train_loader:
img_in, img_out = img_in.to(device), img_out.to(device)
pred = model(img_in)
loss = criterion(pred, img_out)

optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, loss={loss.item():.4f}")
import matplotlib.pyplot as plt

img_in, img_out = train_dataset[0]
with torch.no_grad():
pred = model(img_in.unsqueeze(0).to(device)).cpu().squeeze()

plt.subplot(1,3,1); plt.imshow(img_in.squeeze(), cmap="gray"); plt.title("Input")
plt.subplot(1,3,2); plt.imshow(pred, cmap="gray"); plt.title("Predicted")
plt.subplot(1,3,3); plt.imshow(img_out.squeeze(), cmap="gray"); plt.title("Target")
plt.show()


I wonder if such a project is feasible in Python with machine learning and if so, how?

Continue reading...