import cv2
import numpy as np
from mrcnn.config import Config

class PredictionConfig(Config):
    NAME = "petrol_station"
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    NUM_CLASSES = 1 + 1 
    DETECTION_MIN_CONFIDENCE = 0.9

def visualize_detections(image_np, results):
    r = results[0]
    output_image = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    color = (0, 255, 0) # Green

    for i in range(len(r['rois'])):
        y1, x1, y2, x2 = r['rois'][i]
        score = r['scores'][i]
        mask = r['masks'][:, :, i]

        # Draw Mask
        mask_overlay = output_image.copy()
        for c in range(3):
            mask_overlay[:, :, c] = np.where(mask == 1, color[c], output_image[:, :, c])
        cv2.addWeighted(mask_overlay, 0.5, output_image, 0.5, 0, output_image)

        # Draw Box
        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)

        # Draw Label
        label = f"Petrol Station: {score:.2f}"
        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(output_image, (x1, y1 - 20), (x1 + w, y1), color, -1)
        cv2.putText(output_image, label, (x1, y1 - 5), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
                    
    return output_image