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()