Spaces:
Build error
Build error
File size: 2,057 Bytes
d5779bb |
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 |
from __future__ import print_function
import os
import argparse
import numpy as np
import cv2
from Pytorch_Retinaface.pytorch_retinaface import Pytorch_RetinaFace
parser = argparse.ArgumentParser(description='Retinaface')
parser.add_argument('-m', '--trained_model', default='./weights/mobilenet0.25_Final.pth',
type=str, help='Trained state_dict file path to open')
parser.add_argument('--network', default='mobile0.25', help='Backbone network mobile0.25 or resnet50')
parser.add_argument('--cpu', action="store_true", default=False, help='Use cpu inference')
parser.add_argument('--confidence_threshold', default=0.02, type=float, help='confidence_threshold')
parser.add_argument('--top_k', default=5000, type=int, help='top_k')
parser.add_argument('--nms_threshold', default=0.4, type=float, help='nms_threshold')
parser.add_argument('--keep_top_k', default=750, type=int, help='keep_top_k')
parser.add_argument('-s', '--save_image', action="store_true", default=True, help='show detection results')
parser.add_argument('--vis_thres', default=0.6, type=float, help='visualization_threshold')
args = parser.parse_args()
def main():
rf = Pytorch_RetinaFace()
current_dir = os.path.dirname(os.path.abspath(__file__))
output_dir = os.path.join(current_dir, "./Pytorch_Retinaface/outputs")
image_path = os.path.join(current_dir, "./Pytorch_Retinaface/images/test.webp")
if not os.path.exists(image_path):
raise FileNotFoundError
img_raw = cv2.imread(image_path, cv2.IMREAD_COLOR)
img = np.float32(img_raw)
dets = rf.detect_faces(img)
# Crop and save each detected face
cropped_imgs = rf.center_and_crop_rescale(img_raw, dets)
for index, cropped_img in enumerate(cropped_imgs[0]):
# Save the final image
os.makedirs(output_dir, exist_ok=True)
cv2.imwrite(os.path.join(output_dir, f"cropped_face_{index}.jpg"), cropped_img)
print(f"Saved: cropped_face_{index}.jpg")
if __name__ == '__main__':
main() |