import cv2
import numpy as np
import os

import cv2
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Conv2D, Conv2DTranspose, BatchNormalization, LeakyReLU, Dropout, Input, Concatenate
from tensorflow.keras.models import Model

os.makedirs('shapes/incomplete', exist_ok=True)
os.makedirs('shapes/complete', exist_ok=True)

def draw_shape(image, shape):
    image_size = image.shape[0]
    
    if shape == 'circle':
        center = (image_size // 2, image_size // 2)
        radius = image_size // 4
        cv2.circle(image, center, radius, (255, 255, 255), -1)
    
    elif shape == 'rectangle':
        top_left = (image_size // 4, image_size // 4)
        bottom_right = (3 * image_size // 4, 3 * image_size // 4)
        cv2.rectangle(image, top_left, bottom_right, (255, 255, 255), -1)
    
    elif shape == 'triangle':
        points = np.array([[image_size // 2, image_size // 4], 
                           [image_size // 4, 3 * image_size // 4], 
                           [3 * image_size // 4, 3 * image_size // 4]])
        cv2.drawContours(image, [points], 0, (255, 255, 255), -1)
    
    elif shape == 'polygon':
        points = np.array([[image_size // 2, image_size // 4], 
                           [3 * image_size // 4, image_size // 2], 
                           [image_size // 2, 3 * image_size // 4], 
                           [image_size // 4, image_size // 2]])
        cv2.drawContours(image, [points], 0, (255, 255, 255), -1)
    
    return image

def introduce_incompleteness(image):
    mask = np.zeros_like(image)
    
    start_point = (np.random.randint(0, image.shape[1] // 2), 
                   np.random.randint(0, image.shape[0] // 2))
    end_point = (np.random.randint(image.shape[1] // 2, image.shape[1]), 
                 np.random.randint(image.shape[0] // 2, image.shape[0]))
    
    cv2.rectangle(mask, start_point, end_point, (255, 255, 255), -1)
    
    incomplete_image = cv2.bitwise_and(image, image, mask=cv2.bitwise_not(mask))
    
    return incomplete_image

def generate_dataset(num_samples=1000, image_size=128):
    shapes = ['circle', 'rectangle', 'triangle', 'polygon']
    
    X_train = []
    Y_train = []
    
    for i in range(num_samples):
        shape = np.random.choice(shapes)
        complete_image = np.zeros((image_size, image_size), dtype=np.uint8)
        complete_image = draw_shape(complete_image, shape)
        
        incomplete_image = introduce_incompleteness(complete_image.copy())
        
        X_train.append(incomplete_image)
        Y_train.append(complete_image)
        
        cv2.imwrite(f'shapes/incomplete/{shape}_{i}.png', incomplete_image)
        cv2.imwrite(f'shapes/complete/{shape}_{i}.png', complete_image)
    
    return np.array(X_train), np.array(Y_train)

X_train, Y_train = generate_dataset(num_samples=1000, image_size=128)

np.save('X_train.npy', X_train)
np.save('Y_train.npy', Y_train)

X_train = np.load('./X_train.npy')
Y_train = np.load('./Y_train.npy')

X_train = np.array([cv2.resize(img, (256, 256)) for img in X_train])
Y_train = np.array([cv2.resize(img, (256, 256)) for img in Y_train])

X_train = (X_train.astype(np.float32) - 127.5) / 127.5
Y_train = (Y_train.astype(np.float32) - 127.5) / 127.5

X_train = np.expand_dims(X_train, axis=-1)
Y_train = np.expand_dims(Y_train, axis=-1)

X_train = tf.convert_to_tensor(X_train)
Y_train = tf.convert_to_tensor(Y_train)

def build_generator(input_shape=(256, 256, 1)):
    def downsample(filters, size, apply_batchnorm=True):
        initializer = tf.random_normal_initializer(0., 0.02)
        result = tf.keras.Sequential()
        result.add(Conv2D(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False))
        if apply_batchnorm:
            result.add(BatchNormalization())
        result.add(LeakyReLU())
        return result

    def upsample(filters, size, apply_dropout=False):
        initializer = tf.random_normal_initializer(0., 0.02)
        result = tf.keras.Sequential()
        result.add(Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False))
        result.add(BatchNormalization())
        if apply_dropout:
            result.add(Dropout(0.5))
        result.add(tf.keras.layers.ReLU())
        return result

    inputs = Input(shape=input_shape)

    down_stack = [
        downsample(64, 4, apply_batchnorm=False),
        downsample(128, 4),
        downsample(256, 4),
        downsample(512, 4),
        downsample(512, 4),
        downsample(512, 4),
        downsample(512, 4),
        downsample(512, 4),
    ]

    up_stack = [
        upsample(512, 4, apply_dropout=True),
        upsample(512, 4, apply_dropout=True),
        upsample(512, 4, apply_dropout=True),
        upsample(512, 4),
        upsample(256, 4),
        upsample(128, 4),
        upsample(64, 4),
    ]

    initializer = tf.random_normal_initializer(0., 0.02)
    last = Conv2DTranspose(1, 4, strides=2, padding='same', kernel_initializer=initializer, activation='tanh')

    x = inputs
    skips = []
    for down in down_stack:
        x = down(x)
        skips.append(x)

    skips = reversed(skips[:-1])

    for up, skip in zip(up_stack, skips):
        x = up(x)
        x = Concatenate()([x, skip])

    x = last(x)

    return Model(inputs=inputs, outputs=x)

def build_discriminator(input_shape=(256, 256, 1)):
    initializer = tf.random_normal_initializer(0., 0.02)

    inp = Input(shape=input_shape, name='input_image')
    tar = Input(shape=input_shape, name='target_image')

    x = Concatenate()([inp, tar])

    down1 = Conv2D(64, 4, strides=2, padding='same', kernel_initializer=initializer)(x)
    down1 = LeakyReLU()(down1)

    down2 = Conv2D(128, 4, strides=2, padding='same', kernel_initializer=initializer)(down1)
    down2 = BatchNormalization()(down2)
    down2 = LeakyReLU()(down2)

    down3 = Conv2D(256, 4, strides=2, padding='same', kernel_initializer=initializer)(down2)
    down3 = BatchNormalization()(down3)
    down3 = LeakyReLU()(down3)

    zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)
    conv = Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1)
    batchnorm1 = BatchNormalization()(conv)
    leaky_relu = LeakyReLU()(batchnorm1)

    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)

    last = Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2)

    return Model(inputs=[inp, tar], outputs=last)

def generator_loss(disc_generated_output, gen_output, target):
    LAMBDA = 100
    gan_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)(tf.ones_like(disc_generated_output), disc_generated_output)
    l1_loss = tf.reduce_mean(tf.abs(target - gen_output))
    total_gen_loss = gan_loss + (LAMBDA * l1_loss)
    return total_gen_loss

def discriminator_loss(disc_real_output, disc_generated_output):
    real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)(tf.ones_like(disc_real_output), disc_real_output)
    generated_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)(tf.zeros_like(disc_generated_output), disc_generated_output)
    total_disc_loss = real_loss + generated_loss
    return total_disc_loss

generator = build_generator()
discriminator = build_discriminator()

generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

epochs = 10
batch_size = 32

for epoch in range(epochs):
    for i in range(0, len(X_train), batch_size):
        X_batch = X_train[i:i+batch_size]
        Y_batch = Y_train[i:i+batch_size]
        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
            gen_output = generator(X_batch, training=True)
            disc_real_output = discriminator([X_batch, Y_batch], training=True)
            disc_generated_output = discriminator([X_batch, gen_output], training=True)

            gen_loss = generator_loss(disc_generated_output, gen_output, Y_batch)
            disc_loss = discriminator_loss(disc_real_output, disc_generated_output)

        gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
        gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

        generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
        discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

    print(f'Epoch {epoch+1}, Generator loss: {gen_loss.numpy()}, Discriminator loss: {disc_loss.numpy()}')

generator.save('image_completion_generator.h5')
discriminator.save('image_completion_discriminator.h5')

import cv2
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

generator = tf.keras.models.load_model('./image_completion_generator.h5')
generator.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

def prepare_input_image(image_path):
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  
    image_resized = cv2.resize(image, (256, 256))
    image_normalized = image_resized.astype('float32') / 255.0
    input_image = np.expand_dims(image_normalized, axis=-1)  
    input_image = np.expand_dims(input_image, axis=0)  
    
    return input_image

def generate_completed_image(generator, input_image):
    completed_image = generator.predict(input_image)[0] 
    completed_image = (completed_image > 0.5).astype('uint8') * 255 
    return completed_image

def visualize_images(input_image, completed_image):
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(input_image.squeeze(), cmap='gray')
    plt.title("Input Image")
    plt.axis('off')
    plt.subplot(1, 2, 2)
    plt.imshow(completed_image.squeeze(), cmap='gray')
    plt.title("Completed Image")
    plt.axis('off')
    plt.show()

image_path = 'circle.png'
input_image = prepare_input_image(image_path)
completed_image = generate_completed_image(generator, input_image)
visualize_images(input_image, completed_image)

i made this script to train a model to complete incomplete shapes but it isn't very precise like the accuracy is very low, suggest changes and give updated code