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