File size: 6,238 Bytes
883856e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """
Reusable Inference Module for Text-Conditioned Segmentation
Supports prompt-based image segmentation using trained ResNet18 + UNet model
"""
import torch
import torch.nn as nn
import numpy as np
from PIL import Image, ImageDraw
from torchvision import transforms, models
import os
import cv2
import numpy as np
# ============================================
# CONFIG
# ============================================
DEVICE = "cpu" # RTX 5070 (sm_120) not yet supported
# Prompt → Mode mapping for semantic control
# Although model is not explicitly text-conditioned,
# we use different thresholds based on semantic understanding of the prompt
PROMPT_TO_MODE = {
"segment crack": "crack",
"segment wall crack": "crack",
"segment taping area": "taping",
"segment joint": "taping",
"segment drywall seam": "taping",
}
# Mode-specific thresholds (determined empirically from training)
# The model outputs values in range ~0.4817-0.4874 for most inputs
# We use thresholds very close to the min value to create differentiation
# Cracks: threshold closer to min for better detection
# Taping: threshold slightly higher for more selective detection
MODE_THRESHOLDS = {
"crack": 0.485, # Increased to avoid full-image selection
"taping": 0.486, # Increased for selectivity
}
# ============================================
# MODEL DEFINITION
# ============================================
from model import ResNetSegmentation
# ============================================
# HELPER FUNCTIONS
# ============================================
def draw_rectangles_on_mask(image_pil, mask_binary, thickness=2, color=(0, 255, 0)):
"""
Draw rectangles around detected regions (connected components) on the image.
"""
# Convert to numpy for contour detection
mask_uint8 = np.uint8(mask_binary)
# Find contours using OpenCV
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Copy image for drawing
image_marked = image_pil.copy()
draw = ImageDraw.Draw(image_marked)
# Draw rectangles around contours
for contour in contours:
if len(contour) < 3: # Skip very small contours (points/lines)
continue
# Get bounding rectangle
x, y, w, h = cv2.boundingRect(contour)
# Skip small noise
if w < 10 or h < 10:
continue
# Draw rectangle on PIL Image
draw.rectangle(
[(x, y), (x + w, y + h)],
outline=color,
width=thickness
)
return image_marked
# ============================================
# INFERENCE PIPELINE
# ============================================
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # DeepLab expects normalized input
])
def predict(image_path, prompt, model_path="best_model.pth"):
# Load model
model = ResNetSegmentation().to(DEVICE)
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model file not found: {model_path}")
model.load_state_dict(torch.load(model_path, map_location=DEVICE))
model.eval()
# Load and preprocess image
if not os.path.exists(image_path):
raise FileNotFoundError(f"Image file not found: {image_path}")
image = Image.open(image_path).convert("RGB")
original_size = image.size
# DeepLab expects normalization
img_tensor = transform(image).unsqueeze(0).to(DEVICE)
# Run inference
with torch.no_grad():
logits = model(img_tensor) # Returns logits [B, 1, H, W]
probs = torch.sigmoid(logits) # Convert to [0, 1]
probs_np = probs[0, 0].cpu().numpy()
# Standard Thresholding for a trained model
threshold = 0.5
mask_binary = (probs_np > threshold).astype(np.uint8) * 255
# Resize mask back to original image size
mask_pil = Image.fromarray(mask_binary)
mask_resized = mask_pil.resize(original_size, Image.Resampling.NEAREST)
mask = np.array(mask_resized)
return image, mask
def predict_and_save(image_path, prompt, output_dir="outputs", model_path="best_model.pth"):
"""
Run inference and save the predicted mask as PNG.
Args:
image_path (str): Path to input image
prompt (str): Natural language prompt
output_dir (str): Directory to save output PNG
model_path (str): Path to trained model weights
Returns:
str: Path to saved mask PNG file
"""
# Create output directory if not exists
os.makedirs(output_dir, exist_ok=True)
# Run inference
image, mask = predict(image_path, prompt, model_path)
# Generate output filename
image_name = os.path.splitext(os.path.basename(image_path))[0]
prompt_slug = prompt.replace(" ", "_").lower()
output_filename = f"{image_name}__{prompt_slug}.png"
output_path = os.path.join(output_dir, output_filename)
# Save mask
Image.fromarray(mask).save(output_path)
return output_path
# ============================================
# METRICS
# ============================================
def compute_dice(pred, gt, threshold=0.5):
"""Compute Dice Score"""
pred = (pred > threshold).astype(np.float32)
gt = gt.astype(np.float32)
inter = (pred * gt).sum()
dice = (2 * inter) / (pred.sum() + gt.sum() + 1e-6)
return dice
def compute_iou(pred, gt, threshold=0.5):
"""Compute Intersection over Union"""
pred = (pred > threshold).astype(np.float32)
gt = gt.astype(np.float32)
inter = (pred * gt).sum()
union = pred.sum() + gt.sum() - inter
iou = inter / (union + 1e-6)
return iou
if __name__ == "__main__":
# Example usage
print("Inference module ready. Import and use predict() function.")
print(f"Supported prompts: {list(PROMPT_TO_CLASS.keys())}")
|