Source code for endoreg_db.models.metadata.video_prediction_meta
from typing import Union, TYPE_CHECKING, Optional, List, Tuple
from django.db import models
from django.db.models import Q, CheckConstraint
import logging
from endoreg_db.models.label import LabelSet
from ..label.label_video_segment import (
LabelVideoSegment,
)
from ..utils import find_segments_in_prediction_array
import numpy as np
import pickle
logger = logging.getLogger(__name__)
DEFAULT_WINDOW_SIZE_IN_SECONDS_FOR_RUNNING_MEAN = 1.5
DEFAULT_VIDEO_SEGMENT_LENGTH_THRESHOLD_IN_S = 1.0
if TYPE_CHECKING:
from endoreg_db.models import ModelMeta, InformationSource, Label
from ..prediction.image_classification import ImageClassificationPrediction
from ..media.video.video_file import VideoFile
[docs]
class VideoPredictionMeta(models.Model):
"""
Stores metadata about predictions made by a model for a specific video.
Must be associated with exactly one `VideoFile`.
"""
model_meta = models.ForeignKey("ModelMeta", on_delete=models.CASCADE)
date_created = models.DateTimeField(auto_now_add=True)
date_modified = models.DateTimeField(auto_now=True)
prediction_array = models.BinaryField(blank=True, null=True)
video_file = models.ForeignKey(
"VideoFile",
on_delete=models.CASCADE,
related_name="video_prediction_meta",
null=False,
blank=False,
)
if TYPE_CHECKING:
model_meta: "ModelMeta"
video_file: "VideoFile"
label_video_segments: "models.Manager[LabelVideoSegment]"
class Meta:
constraints = [
models.UniqueConstraint(fields=['model_meta', 'video_file'], name='unique_prediction_per_video_model')
]
indexes = [
models.Index(fields=["model_meta", "video_file"]),
]
[docs]
def get_video(self) -> "VideoFile":
"""Returns the associated VideoFile instance."""
if self.video_file:
return self.video_file
else:
raise ValueError("VideoPredictionMeta is not associated with a VideoFile.")
def __str__(self):
try:
video_obj = self.get_video()
return f"Prediction Meta for Video {video_obj.uuid} - {self.model_meta.name}"
except ValueError:
return f"Prediction Meta {self.pk} (Error: No VideoFile) - {self.model_meta.name}"
except Exception as e:
logger.warning("Error generating string representation for VideoPredictionMeta %s: %s", self.pk, e)
return f"Prediction Meta {self.pk} (Error: {e})"
[docs]
def get_labelset(self) -> Optional["LabelSet"]:
"""Get the labelset associated with the model."""
return self.model_meta.labelset
[docs]
def get_label_list(self) -> list["Label"]:
"""Get the ordered list of labels from the model's labelset."""
labelset = self.get_labelset()
if labelset:
return labelset.get_labels_in_order()
return []
[docs]
def save_prediction_array(self, prediction_array: np.array):
"""
Save the prediction array to the database.
"""
self.prediction_array = pickle.dumps(prediction_array)
self.save(update_fields=['prediction_array', 'date_modified'])
[docs]
def get_prediction_array(self):
"""
Get the prediction array from the database.
"""
if self.prediction_array is None:
return None
else:
try:
return pickle.loads(self.prediction_array)
except (pickle.UnpicklingError, TypeError, EOFError) as e:
logger.error(f"Error unpickling prediction array for {self}: {e}")
return None
[docs]
def calculate_prediction_array(self, window_size_in_seconds: int = None):
"""
Fetches all predictions for the associated video, labelset, and model meta,
applies smoothing, and saves the resulting binary prediction array.
"""
from ..prediction.image_classification import ImageClassificationPrediction
video_obj = self.get_video()
model_meta = self.model_meta
label_list = self.get_label_list()
num_frames = video_obj.frame_count
if num_frames is None or num_frames <= 0:
logger.warning(f"Cannot calculate prediction array for {video_obj} with invalid frame count ({num_frames}).")
return
if not label_list:
logger.warning(f"No labels found for model {model_meta}. Cannot calculate prediction array.")
return
prediction_array = np.zeros((num_frames, len(label_list)))
base_pred_qs = ImageClassificationPrediction.objects.filter(
model_meta=model_meta,
frame__video_file=video_obj
)
for i, label in enumerate(label_list):
predictions = base_pred_qs.filter(label=label).order_by("frame__frame_number").values_list(
"frame__frame_number", "confidence"
)
confidences = np.full(num_frames, 0.5)
found_predictions = False
for frame_num, confidence in predictions:
if 0 <= frame_num < num_frames:
confidences[frame_num] = confidence
found_predictions = True
else:
logger.warning(f"Prediction found for out-of-bounds frame number {frame_num} (max: {num_frames-1}). Skipping.")
if not found_predictions:
logger.warning(f"No predictions found for label '{label.name}' in {video_obj}. Using default confidence.")
smooth_confidences = self.apply_running_mean(
confidences, window_size_in_seconds
)
binary_predictions = smooth_confidences > 0.5
prediction_array[:, i] = binary_predictions
self.save_prediction_array(prediction_array)
logger.info(f"Calculated and saved prediction array for {self}")
[docs]
def apply_running_mean(self, confidence_array, window_size_in_seconds: int = None):
"""
Apply a running mean filter to the confidence array for smoothing, assuming a padding
of 0.5 for the edges.
"""
video_obj = self.get_video()
fps = video_obj.get_fps()
if fps is None or fps <= 0:
logger.warning(f"Invalid FPS ({fps}) for {video_obj}. Cannot apply running mean. Returning original array.")
return confidence_array
if not window_size_in_seconds:
window_size_in_seconds = DEFAULT_WINDOW_SIZE_IN_SECONDS_FOR_RUNNING_MEAN
window_size_in_frames = int(window_size_in_seconds * fps)
window_size_in_frames = max(window_size_in_frames, 1)
window = np.ones(window_size_in_frames) / window_size_in_frames
pad_size = window_size_in_frames // 2
padded_confidences = np.pad(
confidence_array,
(pad_size, pad_size),
"constant",
constant_values=(0.5, 0.5),
)
running_mean = np.convolve(padded_confidences, window, mode="same")
start_index = pad_size
end_index = start_index + len(confidence_array)
running_mean = running_mean[start_index:end_index]
if running_mean.shape != confidence_array.shape:
logger.warning(f"Running mean output shape {running_mean.shape} differs from input {confidence_array.shape}. Check padding/slicing.")
return confidence_array
return running_mean
[docs]
def create_video_segments_for_label(self, segments: List[Tuple[int, int]], label: "Label"):
"""
Creates LabelVideoSegment instances for the given label and segments.
"""
from endoreg_db.models import InformationSource
video_obj = self.get_video()
information_source, _ = InformationSource.objects.get_or_create(name="prediction")
segments_to_create = []
for start_frame, end_frame in segments:
segment_data = {
"start_frame_number": start_frame,
"end_frame_number": end_frame,
"source": information_source,
"label": label,
"prediction_meta": self,
"video_file": video_obj,
}
if not LabelVideoSegment.objects.filter(
video_file=video_obj,
prediction_meta=self,
label=label,
start_frame_number=start_frame,
end_frame_number=end_frame
).exists():
segments_to_create.append(LabelVideoSegment(**segment_data))
if segments_to_create:
LabelVideoSegment.objects.bulk_create(segments_to_create)
logger.info(f"Created {len(segments_to_create)} video segments for label '{label.name}' in {video_obj}.")
else:
logger.info(f"No new video segments needed for label '{label.name}' in {video_obj}.")
[docs]
def create_video_segments(self, segment_length_threshold_in_s: float = None):
"""
Generates LabelVideoSegments based on the stored prediction array.
"""
if not segment_length_threshold_in_s:
segment_length_threshold_in_s = DEFAULT_VIDEO_SEGMENT_LENGTH_THRESHOLD_IN_S
video_obj = self.get_video()
fps = video_obj.get_fps()
if fps is None or fps <= 0:
logger.warning(f"Cannot create video segments for {video_obj} with invalid FPS ({fps}).")
return
min_frame_length = int(segment_length_threshold_in_s * fps)
min_frame_length = max(min_frame_length, 1)
label_list = self.get_label_list()
prediction_array = self.get_prediction_array()
if prediction_array is None:
logger.info(f"Prediction array not found for {self}. Calculating...")
self.calculate_prediction_array()
prediction_array = self.get_prediction_array()
if prediction_array is None:
logger.error(f"Failed to get or calculate prediction array for {self}. Cannot create segments.")
return
if prediction_array.shape[1] != len(label_list):
logger.warning(f"Prediction array shape {prediction_array.shape} incompatible with label list length {len(label_list)} for {self}.")
return
logger.info(f"Creating video segments for {self} (min length: {min_frame_length} frames)...")
for i, label in enumerate(label_list):
binary_predictions = prediction_array[:, i].astype(bool)
segments = find_segments_in_prediction_array(binary_predictions, min_frame_length)
if segments:
self.create_video_segments_for_label(segments, label)
logger.info(f"Finished creating video segments for {self}.")