Spaces:
Running
Running
| import os | |
| import numpy as np | |
| from matplotlib import rcParams | |
| import matplotlib.pyplot as plt | |
| from requests import get | |
| import streamlit as st | |
| import cv2 | |
| from ultralytics import YOLO | |
| import shutil | |
| PREDICTION_PATH = os.path.join('.', 'predictions') | |
| def load_od_model(): | |
| finetuned_model = YOLO('face_detection_best.pt') | |
| return finetuned_model | |
| def inference(input_image_path: str): | |
| finetuned_model = load_od_model() | |
| results = finetuned_model.predict(input_image_path, | |
| show=False, | |
| save=True, | |
| save_crop=False, | |
| imgsz=640, | |
| conf=0.6, | |
| save_txt=True, | |
| project= PREDICTION_PATH, | |
| show_labels=False, | |
| show_conf=False, | |
| line_width=2, | |
| exist_ok=True) | |
| names = finetuned_model.names | |
| nfaces = 0 | |
| for r in results: | |
| for c in r.boxes.cls: | |
| nfaces += 1 | |
| with placeholder.container(): | |
| st.markdown(f"<h5>{nfaces} faces detected.</h5>", unsafe_allow_html=True) | |
| st.image(os.path.join(PREDICTION_PATH, 'predict', 'input.jpg')) | |
| def files_cleanup(path_: str): | |
| if os.path.exists(path_): | |
| os.remove(path_) | |
| shutil.rmtree(PREDICTION_PATH) | |
| # @st.cache_resource | |
| def get_upload_path(): | |
| upload_file_path = os.path.join('.', 'uploads') | |
| if not os.path.exists(upload_file_path): | |
| os.makedirs(upload_file_path) | |
| upload_filename = "input.jpg" | |
| upload_file_path = os.path.join(upload_file_path, upload_filename) | |
| return upload_file_path | |
| def process_input_image(img_url): | |
| upload_file_path = get_upload_path() | |
| headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/102.0.0.0 Safari/537.36'} | |
| r = get(img_url, headers=headers) | |
| arr = np.frombuffer(r.content, np.uint8) | |
| input_image = cv2.imdecode(arr, cv2.IMREAD_UNCHANGED) | |
| input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB) | |
| input_image = cv2.resize(input_image, (640, 640)) | |
| cv2.imwrite(upload_file_path, cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)) | |
| return upload_file_path | |
| try: | |
| st.markdown("<h3>Face Detection</h3>", unsafe_allow_html=True) | |
| desc = '''Dataset used to fine-tune YOLOv8 | |
| can be found <a href="https://universe.roboflow.com/mohamed-traore-2ekkp/face-detection-mik1i/dataset/24" target="_blank"> | |
| here</a>. | |
| ''' | |
| st.markdown(desc, unsafe_allow_html=True) | |
| img_url = st.text_input("Paste the image URL having faces:", "") | |
| placeholder = st.empty() | |
| if img_url: | |
| placeholder.empty() | |
| img_path = process_input_image(img_url) | |
| inference(img_path) | |
| files_cleanup(img_path) | |
| except Exception as e: | |
| st.error(f'An unexpected error occured: \n{e}') |