Spaces:
Sleeping
Sleeping
File size: 1,882 Bytes
cc816ec | 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 | import torch
import torch.nn as nn
import numpy as np
import cv2
from unet import UNet # <-- your model class file
import torch
import torch.nn as nn
from unet import UNet # rename this to match your actual file name
def load_model(model_path, device):
# Load checkpoint
checkpoint = torch.load(model_path, map_location=device)
# Initialize model with same config as training
model = UNet(in_channels=3, out_channels=1).to(device)
# Load weights (important!)
model.load_state_dict(checkpoint["model_state"])
model.eval()
return model
# ----------------------------
# 2. Preprocess Input Image
# ----------------------------
def preprocess(image):
# image = numpy array (H, W, 3) BGR from OpenCV
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # convert to RGB
image = cv2.resize(image, (512, 512)) # SAME SIZE used during training
image = image / 255.0 # normalize to 0-1
image = np.transpose(image, (2, 0, 1)) # (H,W,3) -> (3,H,W)
image = np.expand_dims(image, axis=0) # add batch dimension: (1,3,H,W)
image = torch.tensor(image, dtype=torch.float32)
return image
# ----------------------------
# 3. Postprocess Model Output
# ----------------------------
def postprocess(mask):
mask = torch.sigmoid(mask) # ensure values are 0-1
mask = mask.cpu().detach().numpy()[0, 0] # (1,1,H,W) -> (H,W)
mask = (mask > 0.5).astype(np.uint8) * 255 # threshold + convert to image
return mask
# ----------------------------
# 4. Prediction Function
# ----------------------------
def predict(model, image_path, device):
image = cv2.imread(image_path)
input_tensor = preprocess(image).to(device)
with torch.no_grad():
output = model(input_tensor)
mask = postprocess(output)
return image, mask
|