File size: 5,267 Bytes
caaaccd |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import streamlit as st
import cv2
import numpy as np
from PIL import Image
from io import BytesIO
import base64
#---------------------------------------------------------------------------------------
# The session_state function allows us to initialize and save variables across for across
# session states. This is a valuable feature that enables us to take different actions
# depending on the state of selected variables in the code. If this is not done then
# all variables are reset any time the application state changes (e.g., when a user
# interacts with a widget). For example, the confidence threshold of the slider has
# changed, but we are still working with the same image, we can detect that by
# comparing the current file_uploaded_id (img_file_buffer.id) with the
# previous value (ss.file_uploaded_id) and if they are the same then we know we
# don't need to call the face detection model again. We just simply need to process
# the previous set of detections.
#---------------------------------------------------------------------------------------
# Create application title and file uploader widget.
st.title("OpenCV Deep Learning based Face Detection")
img_file_buffer = st.file_uploader("Choose a file", type=['jpg', 'jpeg', 'png'])
# Initialize session state variables
if 'file_uploaded_name' not in st.session_state:
st.session_state.file_uploaded_name = None
if 'detections' not in st.session_state:
st.session_state.detections = None
# Function for detecting faces in an image.
def detectFaceOpenCVDnn(net, frame):
# Create a blob from the image and apply some pre-processing.
blob = cv2.dnn.blobFromImage(frame, 1.0, (300, 300), [104, 117, 123], False, False)
# Set the blob as input to the model.
net.setInput(blob)
# Get Detections.
detections = net.forward()
return detections
# Function for annotating the image with bounding boxes for each detected face.
def process_detections(frame, detections, conf_threshold=0.5):
bboxes = []
frame_h = frame.shape[0]
frame_w = frame.shape[1]
# Loop over all detections and draw bounding boxes around each face.
for i in range(detections.shape[2]):
confidence = detections[0, 0, i, 2]
if confidence > conf_threshold:
x1 = int(detections[0, 0, i, 3] * frame_w)
y1 = int(detections[0, 0, i, 4] * frame_h)
x2 = int(detections[0, 0, i, 5] * frame_w)
y2 = int(detections[0, 0, i, 6] * frame_h)
bboxes.append([x1, y1, x2, y2])
bb_line_thickness = max(1, int(round(frame_h / 200)))
# Draw bounding boxes around detected faces.
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), bb_line_thickness, cv2.LINE_8)
return frame, bboxes
# Function to load the DNN model.
@st.cache_resource()
def load_model():
modelFile = "res10_300x300_ssd_iter_140000_fp16.caffemodel"
configFile = "deploy.prototxt"
net = cv2.dnn.readNetFromCaffe(configFile, modelFile)
return net
# Function to generate a download link for output file.
def get_image_download_link(img, filename, text):
buffered = BytesIO()
img.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode()
href = f'<a href="data:file/txt;base64,{img_str}" download="{filename}">{text}</a>'
return href
net = load_model()
if img_file_buffer is not None:
# Read the file and convert it to OpenCV Image.
raw_bytes = np.asarray(bytearray(img_file_buffer.read()), dtype=np.uint8)
# Loads image in a BGR channel order.
image = cv2.imdecode(raw_bytes, cv2.IMREAD_COLOR)
file_name = img_file_buffer.name
# Create placeholders to display input and output images.
placeholders = st.columns(2)
# Display Input image in the first placeholder.
placeholders[0].image(image, channels='BGR')
placeholders[0].text("Input Image")
# Create a Slider and get the threshold from the slider.
conf_threshold = st.slider("SET Confidence Threshold", min_value=0.0, max_value=1.0, step=.01, value=0.5)
# Check if the loaded image is "new", if so call the face detection model function.
if file_name != st.session_state.file_uploaded_name:
# Set the file_uploaded_name equal to the name of the file that was just uploaded.
st.session_state.file_uploaded_name = file_name
# Save the detections in the session-state for future use with the current loaded image.
st.session_state.detections = detectFaceOpenCVDnn(net, image)
st.write("New image uploaded, calling the face detection model.")
else:
st.write("Same image used, processing with the previous detections.")
# Process the detections based on the current confidence threshold.
out_image, _ = process_detections(image, st.session_state.detections, conf_threshold=conf_threshold)
# Display Detected faces.
placeholders[1].image(out_image, channels='BGR')
placeholders[1].text("Output Image")
# Convert OpenCV image to PIL.
out_image = Image.fromarray(out_image[:, :, ::-1])
# Create a link for downloading the output file.
st.markdown(get_image_download_link(out_image, "face_output.jpg", 'Download Output Image'), unsafe_allow_html=True) |