core_pipeline.py

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/yolo11n.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