tree_height / core_pipeline.py
nagasurendra's picture
Update core_pipeline.py
544c4d7 verified
raw
history blame contribute delete
992 Bytes
import cv2
import os
import tempfile
import numpy as np
from ultralytics import YOLO
import matplotlib.pyplot as plt
# Load YOLOv8 model (you can replace with your custom tree model)
model = YOLO("./data/best.pt") # Replace with tree-trained model
def extract_frames(video_path, interval=30):
cap = cv2.VideoCapture(video_path)
frames = []
idx = 0
while True:
ret, frame = cap.read()
if not ret:
break
if idx % interval == 0:
frames.append(frame)
idx += 1
cap.release()
return frames
def detect_trees(frame):
results = model(frame)
bboxes = results[0].boxes.xyxy.cpu().numpy()
confs = results[0].boxes.conf.cpu().numpy()
labels = results[0].boxes.cls.cpu().numpy()
return frame, bboxes, confs, labels
def plot_detections(frame, bboxes):
for box in bboxes:
x1, y1, x2, y2 = box.astype(int)
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
return frame