| 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) |
|
|
| for i in range(len(r['rois'])): |
| y1, x1, y2, x2 = r['rois'][i] |
| score = r['scores'][i] |
| mask = r['masks'][:, :, i] |
|
|
| |
| 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) |
|
|
| |
| cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2) |
|
|
| |
| 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 |