nagasurendra commited on
Commit
286a11e
·
verified ·
1 Parent(s): 99b5388

Create core_pipeline.py

Browse files
Files changed (1) hide show
  1. core_pipeline.py +36 -0
core_pipeline.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import os
3
+ import tempfile
4
+ import numpy as np
5
+ from ultralytics import YOLO
6
+ import matplotlib.pyplot as plt
7
+
8
+ # Load YOLOv8 model (you can use your own tree-detection model)
9
+ model = YOLO("yolov8n.pt") # Replace with tree-trained model
10
+
11
+ def extract_frames(video_path, interval=30):
12
+ cap = cv2.VideoCapture(video_path)
13
+ frames = []
14
+ idx = 0
15
+ while True:
16
+ ret, frame = cap.read()
17
+ if not ret:
18
+ break
19
+ if idx % interval == 0:
20
+ frames.append(frame)
21
+ idx += 1
22
+ cap.release()
23
+ return frames
24
+
25
+ def detect_trees(frame):
26
+ results = model(frame)
27
+ bboxes = results[0].boxes.xyxy.cpu().numpy()
28
+ confs = results[0].boxes.conf.cpu().numpy()
29
+ labels = results[0].boxes.cls.cpu().numpy()
30
+ return frame, bboxes, confs, labels
31
+
32
+ def plot_detections(frame, bboxes):
33
+ for box in bboxes:
34
+ x1, y1, x2, y2 = box.astype(int)
35
+ cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
36
+ return frame