final.py

import gradio as gr
from PIL import Image
import os
import cv2
import numpy as np
import tensorflow as tf

W = 512
H = 512

""" Load the model """
# model_path = os.path.join("/content/drive/MyDrive/colab/", "cutted_full.h5")
edge_model = tf.keras.models.load_model("cutted_full.h5")

""" Load the model """
model_path = "finalize.h5"
extraction_model = tf.keras.models.load_model(model_path)

def read_image(path):
    # path = path.decode()
    x = cv2.imread(path, cv2.IMREAD_COLOR)
    first_5_columns = x[500:-400, :50, :]
    last_5_columns = x[500:-400, -50:, :]
    new_image = np.concatenate((first_5_columns, last_5_columns), axis=1)
    x = cv2.resize(new_image, (H, W))
    x = x / 255.0
    x = np.expand_dims(x, axis=0)
    return x


def edger(image_path,model):
    image = read_image(image_path)
    print("completed reading")
    print(image.shape)

    """ Prediction """
    pred = model.predict(image, verbose=0)

    n = np.array(pred)
    n.shape

    pr = (n * 255).astype(np.uint8)
    return pr[1][0]

def reshaping(original_image):
    image = cv2.resize(original_image,(512,788))

    height, width = image.shape[:2] 

    top_rows = 500
    bottom_rows = 400

    empty_row = np.zeros((1, width), dtype=np.uint8)

    for _ in range(top_rows):
        image = np.vstack((empty_row, image))

    for _ in range(bottom_rows):
        image = np.vstack((image, empty_row))

    return image

def background_generator(image_path,model):
    # Load the original image
    # original = cv2.imread('original.jpg')
    # height, width, channel = original.shape
    height, width = 1688,3008

    background_color = (240, 240, 240)
    image = np.full((height, width, 3), background_color, dtype=np.uint8)

    # Edge Game
    ######################################################################

    original_image = edger(image_path,model)

    original_image = reshaping(original_image)

    # original_image = cv2.resize(original_image,(3008,1688))

    print(original_image.shape)

    # Find the biggest edge on the left side
    left_edge = np.max(original_image[:, :5])
    print(left_edge)

    # Find the biggest edge on the right side
    right_edge = np.max(original_image[:, -5:])
    print(right_edge)

    ######################################################################

    # Draw a curved black line resembling where the wall starts
    line_color = (0, 0, 0)
    line_thickness = 30

    # Use the positions of the left and right maximum points as control points
    left_max_row = np.argmax(original_image[:, 5])
    right_max_row = np.argmax(original_image[:, -5])
    print(left_max_row)
    print(right_max_row)

    # Define the curve using control points
    start_point = (0, left_max_row)
    end_point = (width, right_max_row)
    control_point = (width // 2, int(left_max_row-0.18 * height))

    # Get the height and width of the image
    height, width = image.shape[:2]

    # Calculate the midpoint
    midpoint = height // 2

    # Split the image into upper and lower halves
    upper_half = image[:midpoint, :]
    lower_half = image[midpoint:, :]

    # Change the color of the lower half
    lower_half[:] = (240, 240, 240)

    # Draw the straight line in the middle
    completing_line_start = (1 * width // 10, int(0.49 * height))
    completing_line_end = (8 * width // 10, int(0.49 * height))
    completing_line_color = (240, 240, 240)
    completing_line_thickness = 60
    cv2.line(image, completing_line_start, completing_line_end, completing_line_color, completing_line_thickness, cv2.LINE_AA)

    # Generate a set of points along the curve using quadratic Bezier curve formula
    t = np.linspace(0, 1, 100)
    curve_points = [(int((1 - x) ** 2 * start_point[0] + 2 * (1 - x) * x * control_point[0] + x ** 2 * end_point[0]),
                    int((1 - x) ** 2 * start_point[1] + 2 * (1 - x) * x * control_point[1] + x ** 2 * end_point[1]))
                    for x in t]
    
    for i in range(1, len(curve_points)):
        cv2.line(image, curve_points[i - 1], curve_points[i], line_color, line_thickness, cv2.LINE_AA)

    blur_radius = 9
    image = cv2.GaussianBlur(image, (blur_radius, blur_radius), 0)

    # cv2.imwrite("results/background.jpg", image)
    return image

    

def merging(car_path,background_path):
    car = Image.open(car_path)
    background = Image.open(background_path)
    car = car.resize(background.size)
    merged_image = Image.alpha_composite(background.convert('RGBA'), car.convert('RGBA'))
    return merged_image

""" Creating a directory """
def create_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def prediction(image):
    """ Directory for storing files """
    for item in ["joint", "mask", "extracted"]:
        create_dir(f"results/{item}")
    
    name = "input_image"
    image_path = "results/temp.jpg" 

    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(image_path,image)
    x = cv2.resize(image, (W, H))
    x = x / 255.0
    x = np.expand_dims(x, axis=0)

    """ Prediction """
    pred = extraction_model.predict(x, verbose=0)

    line = np.ones((H, 10, 4)) * 255

    pred_list = []
    for item in pred:
        p = item[0] * 255
        p = np.concatenate([p, p, p, p], axis=-1)

        pred_list.append(p)
        pred_list.append(line)

    cat_images = np.concatenate(pred_list, axis=1)

    """ Save final mask """
    image_h, image_w, _ = image.shape

    y0 = pred[0][0]
    y0 = cv2.resize(y0, (image_w, image_h))
    y0 = np.expand_dims(y0, axis=-1)
    ny = np.where(y0 > 0, 1, y0)

    rgb = image[:, :, 0:3]
    alpha = y0 * 255

    final = np.concatenate((rgb.copy(), alpha), axis=2)
    yy = cv2.merge((ny.copy(), ny.copy(), ny.copy(), y0.copy()))
    mask = yy * 255

    line = np.ones((image_h, 10, 4)) * 255

    # cat_images = np.concatenate([mask, line, final], axis=1)

    # Save the final image with alpha channel and the mask image
    final_image_path = f"results/extracted/{name}.png"
    mask_image_path = f"results/mask/{name}.png"

    cv2.imwrite(final_image_path, final)
    cv2.imwrite(mask_image_path, mask)

    # Read both images with IMREAD_UNCHANGED
    final_image = cv2.imread(final_image_path, cv2.IMREAD_UNCHANGED)
    mask_image = cv2.imread(mask_image_path, cv2.IMREAD_UNCHANGED)

    # Convert to RGB color space
    final_image_rgb = cv2.cvtColor(final_image, cv2.COLOR_BGRA2RGBA)
    # mask_image_rgb = cv2.cvtColor(mask_image, cv2.COLOR_BGR2RGB)

    back = background_generator(image_path,edge_model)
    cv2.imwrite("background.jpg",back)
    final = merging(final_image_path,"background.jpg")
    final.save("results/merged.png")
    
    complete = cv2.imread("results/merged.png")

    complete = cv2.cvtColor(complete, cv2.COLOR_RGB2BGR)
    
    return final_image_rgb, mask_image, complete


# Create a Gradio interface with two output components
iface = gr.Interface(fn=prediction, inputs="image", outputs=["image", "image", "image"], title="Image Segmentation with New Background")
iface.launch()