Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import numpy as np
|
3 |
+
from matplotlib import rcParams
|
4 |
+
import matplotlib.pyplot as plt
|
5 |
+
from requests import get
|
6 |
+
import streamlit as st
|
7 |
+
import cv2
|
8 |
+
from ultralytics import YOLO
|
9 |
+
import shutil
|
10 |
+
|
11 |
+
|
12 |
+
PREDICTION_PATH = os.path.join('.', 'predictions')
|
13 |
+
|
14 |
+
@st.cache_resource
|
15 |
+
def load_od_model():
|
16 |
+
finetuned_model = YOLO('face_detection_best.pt')
|
17 |
+
return finetuned_model
|
18 |
+
|
19 |
+
|
20 |
+
def inference(input_image_path: str):
|
21 |
+
finetuned_model = load_od_model()
|
22 |
+
results = finetuned_model.predict(input_image_path,
|
23 |
+
show=False,
|
24 |
+
save=True,
|
25 |
+
save_crop=False,
|
26 |
+
imgsz=640,
|
27 |
+
conf=0.6,
|
28 |
+
save_txt=True,
|
29 |
+
project= PREDICTION_PATH,
|
30 |
+
show_labels=False,
|
31 |
+
show_conf=False,
|
32 |
+
line_width=2,
|
33 |
+
exist_ok=True)
|
34 |
+
|
35 |
+
names = finetuned_model.names
|
36 |
+
nfaces = 0
|
37 |
+
for r in results:
|
38 |
+
for c in r.boxes.cls:
|
39 |
+
nfaces += 1
|
40 |
+
|
41 |
+
with placeholder.container():
|
42 |
+
st.markdown(f"<h5>{nfaces} faces detected.</h5>", unsafe_allow_html=True)
|
43 |
+
st.image(os.path.join(PREDICTION_PATH, 'predict', 'input.jpg'))
|
44 |
+
|
45 |
+
|
46 |
+
def files_cleanup(path_: str):
|
47 |
+
if os.path.exists(path_):
|
48 |
+
os.remove(path_)
|
49 |
+
shutil.rmtree(PREDICTION_PATH)
|
50 |
+
|
51 |
+
|
52 |
+
@st.cache_resource
|
53 |
+
def get_upload_path():
|
54 |
+
upload_file_path = os.path.join('.', 'uploads')
|
55 |
+
if not os.path.exists(upload_file_path):
|
56 |
+
os.makedirs(upload_file_path)
|
57 |
+
upload_filename = "input.jpg"
|
58 |
+
upload_file_path = os.path.join(upload_file_path, upload_filename)
|
59 |
+
return upload_file_path
|
60 |
+
|
61 |
+
|
62 |
+
def process_input_image(img_url):
|
63 |
+
upload_file_path = get_upload_path()
|
64 |
+
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'}
|
65 |
+
r = get(img_url, headers=headers)
|
66 |
+
arr = np.frombuffer(r.content, np.uint8)
|
67 |
+
input_image = cv2.imdecode(arr, cv2.IMREAD_UNCHANGED)
|
68 |
+
input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
|
69 |
+
input_image = cv2.resize(input_image, (640, 640))
|
70 |
+
cv2.imwrite(upload_file_path, cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR))
|
71 |
+
return upload_file_path
|
72 |
+
|
73 |
+
|
74 |
+
try:
|
75 |
+
st.markdown("<h3>Face Detection</h3>", unsafe_allow_html=True)
|
76 |
+
desc = '''Dataset used to fine-tune YOLOv8
|
77 |
+
can be found <a href="https://universe.roboflow.com/mohamed-traore-2ekkp/face-detection-mik1i/dataset/24" target="_blank">
|
78 |
+
here</a>.
|
79 |
+
'''
|
80 |
+
st.markdown(desc, unsafe_allow_html=True)
|
81 |
+
img_url = st.text_input("Paste the image URL having faces:", "")
|
82 |
+
placeholder = st.empty()
|
83 |
+
if img_url:
|
84 |
+
placeholder.empty()
|
85 |
+
img_path = process_input_image(img_url)
|
86 |
+
inference(img_path)
|
87 |
+
files_cleanup(img_path)
|
88 |
+
|
89 |
+
except Exception as e:
|
90 |
+
st.error(f'An unexpected error occured: \n{e}')
|