MLFinalProject/Final_Project_Object_Localization

1import gradio as gr2import numpy as np3from skimage.color import rgb2gray4from skimage.transform import resize5import matplotlib.pyplot as plt6import matplotlib.patches as patches7from joblib import load8import tempfile9import pandas as pd10from sklearn.preprocessing import StandardScaler11import os12 13# Load models14mlp_model = load("mlp_bbox_model.pkl")15model_x = load("linear_model_x.pkl")16model_y = load("linear_model_y.pkl")17model_w = load("linear_model_w.pkl")18model_h = load("linear_model_h.pkl")19 20# Load annotations and fit scaler (for MLP model)21# IMPORTANT: Load the scaler saved during training so inverse_transform uses22# the exact same mean/std. If target_scaler.pkl is missing, fall back to23# fitting on the full CSV (less accurate — save the scaler from training!).24annotation_data = pd.read_csv("image_annotation.csv")25if os.path.exists("target_scaler.pkl"):26    target_scaler = load("target_scaler.pkl")27else:28    target_scaler = StandardScaler()29    target_scaler.fit(annotation_data[['x', 'y', 'width', 'height']])30 31# Preprocess function32def preprocess_image(image, image_size=(64, 64)):33    if image.ndim == 3 and image.shape[2] == 4:  # RGBA -> RGB34        image = image[:, :, :3]35    if image.ndim == 3 and image.shape[2] == 3:36        image = rgb2gray(image)37    image_resized = resize(image, image_size, anti_aliasing=True)38    return image_resized.flatten(), image39 40# Predict and draw function41def predict(image, model_type):42    x_input, original_image = preprocess_image(image)43    x_input = x_input.reshape(1, -1)44 45    if model_type == "MLP":46        y_scaled_pred = mlp_model.predict(x_input)47        y_pred = target_scaler.inverse_transform(y_scaled_pred)[0]48    else:  # Linear Regression49        x_pred = model_x.predict(x_input)[0]50        y_pred_ = model_y.predict(x_input)[0]51        w_pred = model_w.predict(x_input)[0]52        h_pred = model_h.predict(x_input)[0]53        y_pred = [x_pred, y_pred_, w_pred, h_pred]54 55    # Clip width and height56    y_pred[2] = np.clip(y_pred[2], 1, original_image.shape[1])57    y_pred[3] = np.clip(y_pred[3], 1, original_image.shape[0])58 59    # Attempt to get filename (for ground truth matching)60    # NOTE: Gradio passes numpy arrays; filename matching requires the original61    # filename to be known externally. Ground truth display is unavailable here.62    gt_row = pd.DataFrame()63 64    # Draw image and boxes65    fig, ax = plt.subplots()66    ax.imshow(original_image, cmap='gray')67 68    # Predicted box in red69    rect_pred = patches.Rectangle((y_pred[0], y_pred[1]), y_pred[2], y_pred[3],70                                  linewidth=2, edgecolor='red', facecolor='none', label="Prediction")71    ax.add_patch(rect_pred)72 73    # Ground truth box in green74    if not gt_row.empty:75        x_gt = gt_row.iloc[0]['x']76        y_gt = gt_row.iloc[0]['y']77        w_gt = gt_row.iloc[0]['width']78        h_gt = gt_row.iloc[0]['height']79        rect_gt = patches.Rectangle((x_gt, y_gt), w_gt, h_gt,80                                    linewidth=2, edgecolor='green', facecolor='none', label="Ground Truth")81        ax.add_patch(rect_gt)82 83    ax.legend()84    plt.axis('off')85 86    # Save to temporary file87    tmpfile = tempfile.NamedTemporaryFile(suffix=".png", delete=False)88    tmp_path = tmpfile.name89    tmpfile.close()90    plt.savefig(tmp_path, bbox_inches='tight', pad_inches=0)91    plt.close(fig)92 93    return tmp_path94 95# Gradio interface96interface = gr.Interface(97    fn=predict,98    inputs=[99        gr.Image(label="Upload an image", interactive=True, type="numpy"),100        gr.Radio(choices=["MLP", "Linear Regression"], label="Select Model")101    ],102    outputs=gr.Image(type="filepath", label="Predicted vs Ground Truth"),103    title="Object Localization: Predicted vs Ground Truth",104    description="Upload an image and select a model to predict the bounding box. Ground truth is shown in green if available."105)106 107interface.launch()108