6
0
Эх сурвалжийг харах

Resolve "check video options"

Eric Tröbs 3 жил өмнө
parent
commit
2531d375b7

+ 2 - 2
models/fixed_model/Pipeline.py

@@ -28,9 +28,9 @@ class Pipeline(Interface):
         if file.path in result:
             for r in result[file.path]:
                 if r['type'] == 'MediaBoundingBox':
-                    file.add_bounding_box(r['x'], r['y'], r['w'], r['h'], r['label'])
+                    file.add_bounding_box(r['x'], r['y'], r['w'], r['h'], r['label'], r['frame'])
                 if r['type'] == 'MediaImageLabel':
-                    file.set_image_label(r['label'])
+                    file.set_image_label(r['label'], r['frame'])
 
     def fit(self, storage: MediaStorage):
         print('fmv1 fit')

+ 30 - 12
models/haarcascade_frontalface_default/Pipeline.py

@@ -37,9 +37,33 @@ class Pipeline(Interface):
     def execute(self, storage: MediaStorage, file: MediaFile):
         print('hcffdv1 execute')
 
-        # load image, convert to grayscale, scale down
-        image = cv2.imread(file.path)
+        # load file and analyze frames
+        found = False
 
+        if file.type == 'image':
+            found = self.__find(file, cv2.imread(file.path))
+        else:
+            video = cv2.VideoCapture(file.path)
+            index = 0
+
+            ret, image = video.read()
+            while ret:
+                if self.__find(file, image, index):
+                    found = True
+
+                ret, image = video.read()
+                index += 1
+
+            video.release()
+
+        # set file collection
+        if found:
+            file.set_collection('face')
+        else:
+            file.set_collection('none')
+
+    def __find(self, file: MediaFile, image, frame=None):
+        # convert to grayscale, scale down
         gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
         height, width = gray.shape
 
@@ -55,18 +79,12 @@ class Pipeline(Interface):
             minSize=(192, 192)
         )
 
-        # convert faces to result list
-        result = False
-
+        # add faces to results
         for x, y, w, h in faces:
             file.add_bounding_box(x / scale_width,
                                   y / scale_height,
                                   w / scale_width,
-                                  h / scale_height)
-            result = True
+                                  h / scale_height,
+                                  frame=frame)
 
-        # set file collection
-        if result:
-            file.set_collection('face')
-        else:
-            file.set_collection('none')
+        return len(faces) > 0

+ 1 - 0
pycs/interfaces/MediaBoundingBox.py

@@ -8,6 +8,7 @@ class MediaBoundingBox:
         self.w = result.data['w']
         self.h = result.data['h']
         self.label = result.label
+        self.frame = result.data['frame'] if 'frame' in result.data else None
 
     def serialize(self) -> dict:
         return dict({'type': 'bounding-box'}, **self.__dict__)

+ 9 - 3
pycs/interfaces/MediaFile.py

@@ -37,16 +37,22 @@ class MediaFile:
         self.__file.set_collection_by_reference(reference)
         self.__notifications.add(self.__notifications.nm.edit_file, self.__file)
 
-    def set_image_label(self, label: Union[int, MediaLabel]):
+    def set_image_label(self, label: Union[int, MediaLabel], frame: int = None):
         """
         create a labeled-image result
 
         :param label: label identifier
+        :param frame: frame index (only set for videos)
         """
         if label is not None and isinstance(label, MediaLabel):
             label = label.identifier
 
-        created = self.__file.create_result('pipeline', 'labeled-image', label)
+        if frame is not None:
+            data = {'frame': frame}
+        else:
+            data = None
+
+        created = self.__file.create_result('pipeline', 'labeled-image', label, data)
         self.__notifications.add(self.__notifications.nm.create_result, created)
 
     def add_bounding_box(self, x: float, y: float, w: float, h: float,
@@ -59,7 +65,7 @@ class MediaFile:
         :param w: relative width [0, 1]
         :param h: relative height [0, 1]
         :param label: label
-        :param frame: frame index
+        :param frame: frame index (only set for videos)
         """
         result = {
             'x': x,

+ 1 - 0
pycs/interfaces/MediaImageLabel.py

@@ -4,6 +4,7 @@ from pycs.database.Result import Result
 class MediaImageLabel:
     def __init__(self, result: Result):
         self.label = result.label
+        self.frame = result.data['frame'] if 'frame' in result.data else None
 
     def serialize(self) -> dict:
         return dict({'type': 'image-label'}, **self.__dict__)