saniaE
refactored code
26ed134
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