ksvmuralidhar commited on
Commit
a7ee50e
·
verified ·
1 Parent(s): 4d027fb

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +90 -0
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}')