52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
from enum import Enum
|
|
from data.library import ArkiteData, ArkiteDetection, ArkiteRecording
|
|
from cv2util import to_8_bit_image
|
|
from typing import Iterable
|
|
import numpy
|
|
from recording_grouper import RecordingGrouper
|
|
|
|
class VideoType(Enum):
|
|
IR = 0
|
|
DEPTH = 1
|
|
|
|
class VideoLoader:
|
|
def __init__(self, projects_path: str):
|
|
self.projects_path = projects_path
|
|
# Search folders with depth of 1:
|
|
self.data = ArkiteData(self.projects_path, 1)
|
|
|
|
def get_all_recordings(self) -> Iterable[ArkiteRecording]:
|
|
# Loop over projects
|
|
for detection in self.data.detections(self.data.use_cases()):
|
|
for recording in detection.recordings():
|
|
yield recording
|
|
|
|
def get_recordings_grouped(self):
|
|
grouper = RecordingGrouper(self.data)
|
|
return grouper.getRecordingGroups()
|
|
|
|
@staticmethod
|
|
def extract_frames(recording: ArkiteRecording, video_type: VideoType = None) -> Iterable[numpy.ndarray]:
|
|
if video_type == VideoType.IR or video_type == None:
|
|
disp_min = 100
|
|
disp_max = 8000
|
|
with recording.ir() as ir:
|
|
frames = ir.frame_sequence()
|
|
for frame in frames:
|
|
yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)
|
|
elif video_type == VideoType.DEPTH:
|
|
# TODO: check if these values make sense for depth images.
|
|
disp_min = 100
|
|
disp_max = 8000
|
|
with recording.depth() as depth:
|
|
frames = depth.frame_sequence()
|
|
for frame in frames:
|
|
yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)
|
|
else:
|
|
raise NotImplementedError("VideoType \'{}\' not implemented".format(video_type))
|
|
|
|
# Yielding in sub method after "with" statement in top method does not work!
|
|
# @staticmethod
|
|
# def frames_to_8_bit(frame_sequence, disp_min=None, disp_max=None):
|
|
# for frame in frame_sequence:
|
|
# yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max) |