Update app.py

app.py

@@ -37,24 +37,21 @@ class ValidImgDetector:
         mode: str,
         predictor: Predictor,
     ) -> np.ndarray:
-        # input is RGB image, output must be RGB too
         predictor.detector.detector_kwargs["conf"] = score_threshold
         predictor.detector.detector_kwargs["iou"] = iou_threshold
+
         if mode == "Use persons and faces":
             use_persons = True
             disable_faces = False
-
         elif mode == "Use persons only":
             use_persons = True
             disable_faces = True
-
         elif mode == "Use faces only":
             use_persons = False
             disable_faces = False
 
         predictor.age_gender_model.meta.use_persons = use_persons
         predictor.age_gender_model.meta.disable_faces = disable_faces
-        # image = image[:, :, ::-1]  # RGB -> BGR
         detected_objects, out_im = predictor.recognize(image)
         has_child, has_female, has_male = False, False, False
         if len(detected_objects.ages) > 0:
@@ -64,24 +61,41 @@ class ValidImgDetector:
 
         return out_im[:, :, ::-1], has_child, has_female, has_male
 
-    def valid_img(self, img_path):
-        image = cv2.imread(img_path)
-        return self._detect(image, 0.4, 0.7, "Use persons and faces", self.predictor)
+    def valid_video(self, video_path):
+        cap = cv2.VideoCapture(video_path)
+        results = []
 
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
 
-def infer(photo: str):
-    if is_url(photo):
+            # Process each frame
+            out_frame, has_child, has_female, has_male = self._detect(frame, 0.4, 0.7, "Use persons and faces", self.predictor)
+            results.append((out_frame, has_child, has_female, has_male))
+        
+        cap.release()
+        return results
+
+def infer(video_path: str):
+    if is_url(video_path):
         if os.path.exists(TMP_DIR):
             shutil.rmtree(TMP_DIR)
 
-        photo = download_file(photo, f"{TMP_DIR}/download.jpg")
+        video_path = download_file(video_path, f"{TMP_DIR}/download.mp4")
 
     detector = ValidImgDetector()
-    if not photo or not os.path.exists(photo) or imghdr.what(photo) is None:
-        return None, None, None, "Please input the image correctly"
-
-    return detector.valid_img(photo)
-
+    if not video_path or not os.path.exists(video_path):
+        return None, None, None, "Please input the video correctly"
+
+    results = detector.valid_video(video_path)
+    
+    # Process results to only return first frame's result as an example
+    if results:
+        first_frame_result = results[0]
+        return first_frame_result
+    else:
+        return None, None, None, "No frames detected in video."
 
 if __name__ == "__main__":
     with gr.Blocks() as iface:
@@ -89,22 +103,7 @@ if __name__ == "__main__":
         with gr.Tab("Upload Mode"):
             gr.Interface(
                 fn=infer,
-                inputs=gr.Image(label="Upload Photo", type="filepath"),
-                outputs=[
-                    gr.Image(label="Detection Result", type="numpy"),
-                    gr.Textbox(label="Has Child"),
-                    gr.Textbox(label="Has Female"),
-                    gr.Textbox(label="Has Male"),
-                ],
-                examples=get_jpg_files(f"{MODEL_DIR}/examples"),
-                allow_flagging="never",
-                cache_examples=False,
-            )
-
-        with gr.Tab("Online Mode"):
-            gr.Interface(
-                fn=infer,
-                inputs=gr.Textbox(label="Online Picture URL"),
+                inputs=gr.Video(label="Upload Video"),
                 outputs=[
                     gr.Image(label="Detection Result", type="numpy"),
                     gr.Textbox(label="Has Child"),