Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| import pathlib | |
| import os | |
| import huggingface_hub | |
| import numpy as np | |
| import argparse | |
| from dataclasses import dataclass | |
| from mivolo.predictor import Predictor | |
| from PIL import Image | |
| class Cfg: | |
| detector_weights: str | |
| checkpoint: str | |
| device: str = "cpu" | |
| with_persons: bool = True | |
| disable_faces: bool = False | |
| draw: bool = True | |
| def load_models(): | |
| detector_path = huggingface_hub.hf_hub_download('iitolstykh/demo_yolov8_detector', | |
| 'yolov8x_person_face.pt') | |
| age_gender_path_v1 = 'age_estimator/MiVOLO-main/models/model_imdb_cross_person_4.22_99.46.pth.tar' | |
| predictor_cfg_v1 = Cfg(detector_path, age_gender_path_v1) | |
| predictor_v1 = Predictor(predictor_cfg_v1) | |
| return predictor_v1 | |
| def detect(image: np.ndarray, score_threshold: float, iou_threshold: float, mode: str, predictor: Predictor) -> np.ndarray: | |
| 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 for OpenCV | |
| detected_objects, out_im = predictor.recognize(image) | |
| return out_im[:, :, ::-1] # BGR -> RGB | |
| def load_image(image_path: str): | |
| image = Image.open(image_path) | |
| image_np = np.array(image) | |
| return image_np | |
| def main(args): | |
| # Load models | |
| predictor_v1 = load_models() | |
| # Set parameters from args | |
| score_threshold = args.score_threshold | |
| iou_threshold = args.iou_threshold | |
| mode = args.mode | |
| # Load and process image | |
| image_np = load_image(args.image_path) | |
| # Predict with model | |
| result = detect(image_np, score_threshold, iou_threshold, mode, predictor_v1) | |
| output_image = Image.fromarray(result) | |
| output_image.save(args.output_path) | |
| output_image.show() | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description='Object Detection with YOLOv8 and Age/Gender Prediction') | |
| parser.add_argument('--image_path', type=str, required=True, help='Path to the input image') | |
| parser.add_argument('--output_path', type=str, default='output_image.jpg', help='Path to save the output image') | |
| parser.add_argument('--score_threshold', type=float, default=0.4, help='Score threshold for detection') | |
| parser.add_argument('--iou_threshold', type=float, default=0.7, help='IoU threshold for detection') | |
| parser.add_argument('--mode', type=str, choices=["Use persons and faces", "Use persons only", "Use faces only"], | |
| default="Use persons and faces", help='Detection mode') | |
| args = parser.parse_args() | |
| main(args) | |