Spaces:
Sleeping
Sleeping
| 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 |