| import gradio as gr | |
| import math | |
| import os | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import segmentation_models_pytorch as smp | |
| def pad_to_divisible(img, div=32): | |
| h, w, _ = img.shape | |
| new_h = math.ceil(h / div) * div | |
| new_w = math.ceil(w / div) * div | |
| pad_bottom = new_h - h | |
| pad_right = new_w - w | |
| padded = cv2.copyMakeBorder(img, 0, pad_bottom, 0, pad_right, cv2.BORDER_CONSTANT, value=[0, 0, 0]) | |
| return padded | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("Using device:", device) | |
| model_path = "best_unet_model_complete.pth" | |
| if os.path.exists(model_path): | |
| loaded_model = torch.load(model_path, map_location=device) | |
| loaded_model.eval() | |
| print("Loaded complete model from", model_path) | |
| else: | |
| raise FileNotFoundError(f"Model file not found at {model_path}") | |
| def predict(image): | |
| """ | |
| Takes an input image (as a NumPy array), pads it, performs inference, | |
| and returns: | |
| - the padded input image, | |
| - the predicted mask (grayscale), and | |
| - the original image with a red overlay on the predicted regions. | |
| """ | |
| if image.shape[2] == 4: | |
| image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) | |
| external_image_rgb = image.copy() | |
| external_image_padded = pad_to_divisible(external_image_rgb, div=32) | |
| external_image_norm = external_image_padded.astype(np.float32) / 255.0 | |
| external_tensor = torch.from_numpy(external_image_norm).permute(2, 0, 1) | |
| external_tensor = external_tensor.unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| output = loaded_model(external_tensor) | |
| pred_mask = (torch.sigmoid(output) > 0.3).float() | |
| pred_mask_np = pred_mask.cpu().squeeze().numpy() | |
| overlay_image = external_image_padded.astype(np.float32) | |
| mask_bool = pred_mask_np > 0.5 | |
| red_color = np.array([255, 0, 0], dtype=np.float32) | |
| alpha = 0.5 | |
| overlay_image[mask_bool] = (1 - alpha) * overlay_image[mask_bool] + alpha * red_color | |
| overlay_image = np.clip(overlay_image, 0, 255).astype(np.uint8) | |
| return external_image_padded, pred_mask_np, overlay_image | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="numpy", label="Upload Image"), | |
| outputs=[ | |
| gr.Image(label="Padded Image"), | |
| gr.Image(label="Predicted Mask"), | |
| gr.Image(label="Overlay (Predicted Regions)") | |
| ], | |
| title="Wrinkle Segmentation", | |
| description="Upload an image to see wrinkle segmentation. The app displays the padded image, the predicted mask, and an overlay of the predicted regions in red." | |
| ) | |
| demo.launch() |