conan/container.py

152 lines
5.1 KiB
Python
Raw Permalink Normal View History

2021-10-17 12:49:49 +02:00
import numpy as np
import pandas as pd
import pickle as pkl
import os.path
import cv2
from PyQt5 import QtCore, QtGui
class DataContainer():
def __init__(self, movie_fileName, data_fileName):
self.movie_fileName = movie_fileName
self.data_fileName = data_fileName
self.frameCount = 0
self.fps = 0
self.frameSize = [0, 0]
self.dataGaze = None
self.dataMovement = None
self.number_ids = None
self.dataRTGene = None
self.image = None
def readData(self, verbose=False, rtgene=False):
if (verbose):
print("## Start Reading Data")
# Read Video Data
f = self.movie_fileName
print(f)
if os.path.isfile(f):
cap = cv2.VideoCapture(f)
ret, frame = cap.read()
self.image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
h, w, ch = self.image.shape
bytesPerLine = ch * w
convertToQtFormat = QtGui.QImage(self.image.data, w, h, bytesPerLine, QtGui.QImage.Format_RGB888)
self.fps = cap.get(cv2.CAP_PROP_FPS)
self.frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float
self.frameSize = [width, height]
print(self.fps, self.frameCount)
if (verbose):
print("Video frameCount %i" % self.frameCount)
duration = self.frameCount / self.fps
minutes = int(duration / 60)
seconds = duration % 60
print('Video duration (M:S) = ' + str(minutes) + ':' + str(seconds))
else:
print("WARNING: no video avaibale.")
# read data file
with open(self.data_fileName, 'rb') as f:
data = pkl.load(f)
# Read RT-Gene Gaze
if rtgene:
f = "%s_Gaze.pkl" % self.dataPath[0]
df_gt = pd.read_pickle('exampledata/G2_VID1_GroundTruth.pkl')
df_gt = df_gt[['Frame', 'ID0_target_spher', 'ID1_target_spher']]
if os.path.isfile(f):
df = pd.read_pickle(f)
self.number_ids = len(df.PId.unique())
self.dataRTGene = df.pivot(index='Frame', columns="PId", values=["GazeTheta", "GazePhi", "HeadCenter",
"HeadPoseYaw", "HeadPoseTheta",
"Phi", "Theta"])
lst = []
for label in ["GazeTheta", "GazePhi", "Head", "HeadPoseYaw", "HeadPoseTheta", "Phi", "Theta"]:
for head_id in range(self.number_ids):
lst.append("ID%i_%s" % (head_id, label))
self.dataRTGene.columns = lst
self.dataRTGene = self.dataRTGene.reset_index()
self.dataRTGene = pd.merge(self.dataRTGene, df_gt, on=['Frame'])
self.dataRTGene = self.dataRTGene.rename(
columns={'ID0_target_spher': 'ID1_target', 'ID1_target_spher': 'ID0_target'},
errors='raise')
print('Detected %i IDs in video' % self.number_ids)
if (verbose):
print("Gaze sample count %i" % len(self.dataRTGene))
else:
print("WARNING: no RT-Gene data avaibale.")
# Read Gaze Data
if "HeadPose" in data:
self.dataGaze = data["HeadPose"]
self.number_ids = len([col for col in self.dataGaze.columns if 'head' in col])
print('Detected %i IDs in video' % self.number_ids)
if (verbose):
print("Gaze sample count %i" % len(self.dataGaze))
else:
print("WARNING: no gaze data avaibale.")
"""
# Read OpenPose Data
f = self.__get_filename_with_substring("OpenPose")
if os.path.isfile(f):
self.dataPose = pd.read_pickle(f)
if (verbose):
print("Pose sample count %i" % len(self.dataPose))
else:
print("WARNING: no pose data avaibale.")
"""
# Read Movement Data
if "BodyMovement" in data:
self.dataMovement = data["BodyMovement"]
if verbose:
print('Body movement sample count %i' % len(self.dataMovement))
else:
print('WARNING: no body movement data available.')
def getFrameCount(self):
return self.frameCount
def getFrameSize(self):
return self.frameSize
def getFPS(self):
return self.fps
def getVideo(self):
return self.movie_fileName
def getGazeData(self):
return self.dataGaze
def getFrame(self, frameIdx):
return frameIdx
def getFrameCurrent(self):
return 1
def getNumberIDs(self):
return self.number_ids
def getMovementData(self):
return self.dataMovement
def getRTGeneData(self):
return self.dataRTGene
def getImage(self):
return self.image