FireDetection / firedetection.py
bahakizil's picture
Upload 4 files
1afa0e4 verified
raw
history blame
2.43 kB
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