46dea3304f
This commit integrates the backend API for fetching and updating report data. It also includes a normalization function to handle data consistency between the API and local storage. Co-authored-by: anthonymuncher <anthonymuncher@gmail.com>
42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
# Load your trained YOLOv12 model
|
|
model = YOLO("checkpoints/pothole_detector/weights/best.pt") # Path to your trained weights
|
|
|
|
# Define severity thresholds (you can adjust these based on your dataset)
|
|
def classify_severity(box, image_height):
|
|
x1, y1, x2, y2 = box
|
|
area = (x2 - x1) * (y2 - y1)
|
|
if area > 50000 or y2 > image_height * 0.75:
|
|
return "High"
|
|
elif area > 20000 or y2 > image_height * 0.5:
|
|
return "Medium"
|
|
else:
|
|
return "Low"
|
|
|
|
# Draw bounding boxes with severity
|
|
def draw_boxes_and_severity(image, results):
|
|
for r in results: # iterate over Results objects
|
|
for box in r.boxes.xyxy: # xyxy format
|
|
x1, y1, x2, y2 = map(int, box.cpu().numpy())
|
|
conf = float(r.boxes.conf[0]) if hasattr(r.boxes, "conf") else 0.0
|
|
severity = classify_severity((x1, y1, x2, y2), image.shape[0])
|
|
color = (0, 255, 0) if severity == "Low" else (0, 255, 255) if severity == "Medium" else (0, 0, 255)
|
|
cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
|
|
cv2.putText(image, f"{severity} ({conf:.2f})", (x1, y1 - 10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
|
|
return image
|
|
|
|
# Detect potholes in an image
|
|
def detect_potholes(image_path, output_path="output.jpg"):
|
|
image = cv2.imread(image_path)
|
|
results = model(image) # Run inference
|
|
image = draw_boxes_and_severity(image, results)
|
|
cv2.imwrite(output_path, image)
|
|
print(f"Output saved to {output_path}")
|
|
|
|
# Example usage
|
|
if __name__ == "__main__":
|
|
detect_potholes(r"images\pothole_1.jpg")
|