Spaces:
Running
Running
File size: 2,425 Bytes
1afa0e4 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
from collections import defaultdict
import cv2
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors
def detect_fire_in_video(
input_video_path: str,
output_video_path: str,
model_path: str,
device: str = "cpu"
):
"""
Belirtilen videoda YOLO modelini kullanarak yangın tespiti yapar ve
etiketlenmiş çıktıyı output_video_path'e kaydeder.
:param input_video_path: Girdi videosunun dosya yolu
:param output_video_path: İşlenen (annotate edilmiş) videonun kaydedileceği dosya yolu
:param model_path: YOLO model dosyasının (ör. .pt uzantılı) dosya yolu
:param device: İşlem için kullanılacak donanım ('cpu', 'cuda', 'mps' vb.)
:return: İşlenen video yolunu döndürür.
"""
# Takip geçmişi için defaultdict oluştur
track_history = defaultdict(lambda: [])
# YOLO modelini yükle
model = YOLO(model_path, device=device)
# Video dosyasını aç
cap = cv2.VideoCapture(input_video_path)
# Video özelliklerini al
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH,
cv2.CAP_PROP_FRAME_HEIGHT,
cv2.CAP_PROP_FPS))
# Çıktı videosu kaydı için VideoWriter ayarla
out = cv2.VideoWriter(
output_video_path,
cv2.VideoWriter_fourcc(*"MJPG"),
fps,
(w, h)
)
while True:
ret, im0 = cap.read()
if not ret:
print("Video bitti veya çerçeve okunamadı.")
break
# Annotator objesi ile çizim yapma
annotator = Annotator(im0, line_width=2)
# model.track ile takip (tracking) sonuçlarını al
results = model.track(im0, persist=True)
# Track ID'ler ve maskeler varsa işleme al
if results[0].boxes.id is not None and results[0].masks is not None:
masks = results[0].masks.xy
track_ids = results[0].boxes.id.int().cpu().tolist()
for mask, track_id in zip(masks, track_ids):
annotator.seg_bbox(
mask=mask,
mask_color=colors(int(track_id), True),
label=str(track_id)
)
# Annotate edilmiş frame'i kaydet
out.write(im0)
out.release()
cap.release()
cv2.destroyAllWindows()
return output_video_path
|