Spaces:
Running
Running
File size: 2,645 Bytes
29dba9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from collections import defaultdict
from typing import Dict, Generator, List, Optional, Tuple
import cv2
import numpy as np
import tqdm
from mivolo.model.mi_volo import MiVOLO
from mivolo.model.yolo_detector import Detector
from mivolo.structures import AGE_GENDER_TYPE, PersonAndFaceResult
class Predictor:
def __init__(self, config, verbose: bool = False):
self.detector = Detector(config.detector_weights, config.device, verbose=verbose)
self.age_gender_model = MiVOLO(
config.checkpoint,
config.device,
half=True,
use_persons=config.with_persons,
disable_faces=config.disable_faces,
verbose=verbose,
)
self.draw = config.draw
def recognize(self, image: np.ndarray) -> Tuple[PersonAndFaceResult, Optional[np.ndarray]]:
detected_objects: PersonAndFaceResult = self.detector.predict(image)
self.age_gender_model.predict(image, detected_objects)
out_im = None
if self.draw:
# plot results on image
out_im = detected_objects.plot()
return detected_objects, out_im
def recognize_video(self, source: str) -> Generator:
video_capture = cv2.VideoCapture(source)
if not video_capture.isOpened():
raise ValueError(f"Failed to open video source {source}")
detected_objects_history: Dict[int, List[AGE_GENDER_TYPE]] = defaultdict(list)
total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
for _ in tqdm.tqdm(range(total_frames)):
ret, frame = video_capture.read()
if not ret:
break
detected_objects: PersonAndFaceResult = self.detector.track(frame)
self.age_gender_model.predict(frame, detected_objects)
current_frame_objs = detected_objects.get_results_for_tracking()
cur_persons: Dict[int, AGE_GENDER_TYPE] = current_frame_objs[0]
cur_faces: Dict[int, AGE_GENDER_TYPE] = current_frame_objs[1]
# add tr_persons and tr_faces to history
for guid, data in cur_persons.items():
# not useful for tracking :)
if None not in data:
detected_objects_history[guid].append(data)
for guid, data in cur_faces.items():
if None not in data:
detected_objects_history[guid].append(data)
detected_objects.set_tracked_age_gender(detected_objects_history)
if self.draw:
frame = detected_objects.plot()
yield detected_objects_history, frame
|