import os
import cv2
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog

# Step 1: Configure the model for inference
def configure_model(output_dir, num_classes):
    cfg = get_cfg()
    cfg.MODEL.DEVICE = "cpu"
    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))

    cfg.MODEL.WEIGHTS = os.path.join(output_dir, "model_final.pth")  # Load trained weights
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2  # Confidence threshold
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes  # Update based on your dataset
    # cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"  # Use GPU if available
    return cfg

# Step 2: Run inference and visualize results
def test_model(cfg, image_path, output_path):
    predictor = DefaultPredictor(cfg)
    image = cv2.imread(image_path)  # Load the image
    outputs = predictor(image)  # Run inference
    print(cfg.DATASETS.TRAIN[0])
    # Visualize predictions
    v = Visualizer(image[:, :, ::-1], MetadataCatalog.get("floorplan_train"), scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    
    # Save or display the result
    cv2.imwrite(output_path, out.get_image()[:, :, ::-1])
    print(f"Prediction saved to {output_path}")

# Example usage
if __name__ == "__main__":
    output_dir = "./output"  # Path to the directory where model_final.pth is saved
    image_path = "./1.jpg"  # Path to the image for testing
    output_image = "output_prediction-1.jpg"  # Path to save the prediction result

    # Update with your number of classes
    category = [
    "door1",
    "window2",
    "window2",
    "window1",
    "door1",
    "door1",
    "table1",
    "armchair",
    "table1",
    "tub",
    "sink4",
    "sink3",
    "table1",
    "window2",
    "window2",
    "table2",
    "bed",
    "table1",
    "door1",
    "sofa2",
    "table1",
    "armchair",
    "sink3",
    "armchair",
    "armchair",
    "armchair",
    "table3"]  # Replace with your dataset classes
    categories = list(map(lambda x: {"id": category.index(x), "name": x}, dict.fromkeys(category)))
    num_classes = len(categories)

    # Configure and test the model
    cfg = configure_model(output_dir, num_classes)
    test_model(cfg, image_path, output_image)