Upload 3 files

cutted_full.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02ce1cdf9f00d2d984231406815842f9335424c154abb0c5f59701a6e185f882
+size 530371808

final.py

@@ -0,0 +1,234 @@
+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()

finalize.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eaeda232f5fe6f2f59df7ba6ffd73d198278f38288af1e50ca430f2fe12a777d
+size 530371464