Spaces:
Sleeping
Sleeping
File size: 992 Bytes
286a11e cf8d1d5 544c4d7 286a11e cf8d1d5 286a11e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
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
|