import argparse
import cv2
import numpy as np
import torch

def preprocess(image_path, size=(256, 256)):
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"Could not read image: {image_path}")

    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, size)
    image = image.astype(np.float32) / 255.0
    image = np.transpose(image, (2, 0, 1))

    return torch.from_numpy(image).unsqueeze(0)

def save_mask(mask, output_path):
    mask = mask.squeeze()
    if torch.is_tensor(mask):
        mask = mask.cpu().numpy()

    mask = (mask > 0.5).astype(np.uint8) * 255
    cv2.imwrite(output_path, mask)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--model", default="road_model_Deployment.pt")
    parser.add_argument("--output", default="prediction_mask.png")
    args = parser.parse_args()

    model = torch.jit.load(args.model)
    model.eval()

    image_tensor = preprocess(args.image)

    with torch.no_grad():
        prediction = model(image_tensor)

    save_mask(prediction, args.output)
    print(f"Saved mask to {args.output}")

if __name__ == "__main__":
    main()