Spaces:
Runtime error
Runtime error
#create a Streamlit app using info from image_demo.py | |
import cv2 | |
import time | |
import argparse | |
import os | |
import torch | |
import posenet | |
import tempfile | |
from posenet.utils import * | |
import streamlit as st | |
from posenet.decode_multi import * | |
from visualizers import * | |
from ground_truth_dataloop import * | |
import cv2 | |
import time | |
import argparse | |
import os | |
import torch | |
import posenet | |
import streamlit as st | |
from posenet.decode_multi import * | |
from visualizers import * | |
from ground_truth_dataloop import * | |
st.title('PoseNet Image Analyzer') | |
def process_frame(frame, scale_factor, output_stride): | |
input_image, draw_image, output_scale = process_input(frame, scale_factor=scale_factor, output_stride=output_stride) | |
return input_image, draw_image, output_scale | |
def load_model(model): | |
model = posenet.load_model(model) | |
model = model.cuda() | |
return model | |
def main(): | |
MAX_FILE_SIZE = 20 * 1024 * 1024 # 20 MB | |
model_number = st.sidebar.selectbox('Model', [101, 100, 75, 50]) | |
scale_factor = 1.0 | |
output_stride = st.sidebar.selectbox('Output Stride', [8, 16, 32, 64]) | |
min_pose_score = st.sidebar.number_input("Minimum Pose Score", min_value=0.000, max_value=1.000, value=0.10, step=0.001) | |
st.sidebar.markdown(f'<p style="color:grey; font-size: 12px">The current number is {min_pose_score:.3f}</p>', unsafe_allow_html=True) | |
min_part_score = st.sidebar.number_input("Minimum Part Score", min_value=0.000, max_value=1.000, value=0.010, step=0.001) | |
st.sidebar.markdown(f'<p style="color:grey; font-size:12px">The current number is {min_part_score:.3f}</p>', unsafe_allow_html=True) | |
model = load_model(model_number) | |
output_stride = model.output_stride | |
output_dir = st.sidebar.text_input('Output Directory', './output') | |
option = st.selectbox('Choose an option', ['Upload Image', 'Upload Video', 'Try existing image']) | |
if option == 'Upload Video': | |
video_display_mode = st.selectbox("Video Display Mode", ['Frame by Frame', 'Entire Video']) | |
uploaded_video = st.file_uploader("Upload a video (mp4, mov, avi)", type=['mp4', 'mov', 'avi']) | |
if uploaded_video is not None: | |
tfile = tempfile.NamedTemporaryFile(delete=False) | |
tfile.write(uploaded_video.read()) | |
vidcap = cv2.VideoCapture(tfile.name) | |
success, image = vidcap.read() | |
frames = [] | |
frame_count = 0 | |
while success: | |
input_image, draw_image, output_scale = process_frame(image, scale_factor, output_stride) | |
pose_scores, keypoint_scores, keypoint_coords = run_model(input_image, model, output_stride, output_scale) | |
result_image = posenet.draw_skel_and_kp( | |
draw_image, pose_scores, keypoint_scores, keypoint_coords, | |
min_pose_score=min_pose_score, min_part_score=min_part_score) | |
result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB) | |
# result_image = print_frame(draw_image, pose_scores, keypoint_scores, keypoint_coords, output_dir, min_part_score=min_part_score, min_pose_score=min_pose_score) | |
if result_image is not None: | |
frames.append(result_image) | |
success, image = vidcap.read() | |
frame_count += 1 | |
if video_display_mode == 'Frame by Frame': | |
st.image(result_image, caption=f'Frame {frame_count}', use_column_width=True) | |
# Progress bar | |
progress_bar = st.progress(0) | |
# Write the output video | |
output_file = 'output.mp4' | |
height, width, layers = frames[0].shape | |
size = (width,height) | |
output_file_path = os.path.join(output_dir, output_file) | |
out = cv2.VideoWriter(output_file_path, cv2.VideoWriter_fourcc(*'mp4v'), 15, size) | |
for i in range(len(frames)): | |
progress_percentage = i / len(frames) | |
progress_bar.progress(progress_percentage) | |
out.write(cv2.cvtColor(frames[i], cv2.COLOR_RGB2BGR)) | |
out.release() | |
# Display the processed video | |
if video_display_mode == 'Entire Video': | |
with open(output_file_path, "rb") as file: | |
bytes_data = file.read() | |
st.download_button( | |
label="Download video", | |
data=bytes_data, | |
file_name=output_file, | |
mime="video/mp4", | |
) | |
# video_file = open(output_file_path, 'rb') | |
# st.write(f"Output file path: {output_file_path}") | |
# video_bytes = video_file.read() | |
# st.video(video_bytes) | |
# try: | |
# st.video(bytes_data, format="video/mp4", start_time=0) | |
# # st.write(f"Output file path: {output_file_path}") | |
# # st.video('./output/output.mp4', format="video/mp4", start_time=0) | |
# except Exception as e: | |
# st.write("Error: ", str(e)) | |
if frames: | |
frame_idx = st.slider('Choose a frame', 0, len(frames) - 1, 0) | |
input_image, draw_image, output_scale = process_frame(frames[frame_idx], scale_factor, output_stride) | |
pose_scores, keypoint_scores, keypoint_coords = run_model(input_image, model, output_stride, output_scale) | |
pose_data = { | |
'pose_scores': pose_scores.tolist(), | |
'keypoint_scores': keypoint_scores.tolist(), | |
'keypoint_coords': keypoint_coords.tolist() | |
} | |
st.image(draw_image, caption=f'Frame {frame_idx + 1}', use_column_width=True) | |
st.write(pose_data) | |
progress_bar.progress(1.0) | |
elif option == 'Upload Image': | |
image_file = st.file_uploader("Upload Image (Max 10MB)", type=['png', 'jpg', 'jpeg']) | |
if image_file is not None: | |
if image_file.size > MAX_FILE_SIZE: | |
st.error("File size exceeds the 10MB limit. Please upload a smaller file.") | |
file_bytes = np.asarray(bytearray(image_file.read()), dtype=np.uint8) | |
input_image = cv2.imdecode(file_bytes, 1) | |
filename = image_file.name | |
# Crop the image here as needed | |
# input_image = input_image[y:y+h, x:x+w] | |
input_image, source_image, output_scale = process_input( | |
input_image, scale_factor, output_stride) | |
pose_scores, keypoint_scores, keypoint_coords = run_model(input_image, model, output_stride, output_scale) | |
print_frame(source_image, pose_scores, keypoint_scores, keypoint_coords, output_dir, filename=filename, min_part_score=min_part_score, min_pose_score=min_pose_score) | |
else: | |
st.sidebar.warning("Please upload an image.") | |
elif option == 'Try existing image': | |
image_dir = st.sidebar.text_input('Image Directory', './images_train') | |
if output_dir: | |
if not os.path.exists(output_dir): | |
os.makedirs(output_dir) | |
filenames = [f.path for f in os.scandir(image_dir) if f.is_file() and f.path.endswith(('.png', '.jpg'))] | |
if filenames: | |
selected_image = st.sidebar.selectbox('Choose an image', filenames) | |
input_image, draw_image, output_scale = posenet.read_imgfile( | |
selected_image, scale_factor=scale_factor, output_stride=output_stride) | |
filename = os.path.basename(selected_image) | |
result_image, pose_scores, keypoint_scores, keypoint_coords = run_model(input_image, draw_image, model, output_stride, output_scale) | |
print_frame(result_image, pose_scores, keypoint_scores, keypoint_coords, output_dir, filename=selected_image, min_part_score=min_part_score, min_pose_score=min_pose_score) | |
else: | |
st.sidebar.warning("No images found in directory.") | |
#same as utils.py _process_input | |
def process_input(source_img, scale_factor=1.0, output_stride=16): | |
target_width, target_height = posenet.valid_resolution( | |
source_img.shape[1] * scale_factor, source_img.shape[0] * scale_factor, output_stride=output_stride) | |
scale = np.array([source_img.shape[0] / target_height, source_img.shape[1] / target_width]) | |
input_img = cv2.resize(source_img, (target_width, target_height), interpolation=cv2.INTER_LINEAR) | |
input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB).astype(np.float32) | |
input_img = input_img * (2.0 / 255.0) - 1.0 | |
input_img = input_img.transpose((2, 0, 1)).reshape(1, 3, target_height, target_width) | |
return input_img, source_img, scale | |
def run_model(input_image, model, output_stride, output_scale): | |
with torch.no_grad(): | |
input_image = torch.Tensor(input_image).cuda() | |
heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = model(input_image) | |
# st.text("model heatmaps_result shape: {}".format(heatmaps_result.shape)) | |
# st.text("model offsets_result shape: {}".format(offsets_result.shape)) | |
pose_scores, keypoint_scores, keypoint_coords, pose_offsets = posenet.decode_multi.decode_multiple_poses( | |
heatmaps_result.squeeze(0), | |
offsets_result.squeeze(0), | |
displacement_fwd_result.squeeze(0), | |
displacement_bwd_result.squeeze(0), | |
output_stride=output_stride, | |
max_pose_detections=10, | |
min_pose_score=0.0) | |
# st.text("decoded pose_scores shape: {}".format(pose_scores.shape)) | |
# st.text("decoded pose_offsets shape: {}".format(pose_offsets.shape)) | |
keypoint_coords *= output_scale | |
# Convert BGR to RGB | |
return pose_scores, keypoint_scores, keypoint_coords | |
def print_frame(draw_image, pose_scores, keypoint_scores, keypoint_coords, output_dir, filename=None, min_part_score=0.01, min_pose_score=0.1): | |
if output_dir: | |
draw_image = posenet.draw_skel_and_kp( | |
draw_image, pose_scores, keypoint_scores, keypoint_coords, | |
min_pose_score=min_pose_score, min_part_score=min_part_score) | |
draw_image = cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB) | |
if filename: | |
cv2.imwrite(os.path.join(output_dir, filename), draw_image) | |
else: | |
cv2.imwrite(os.path.join(output_dir, "output.png"), draw_image) | |
st.image(draw_image, caption='PoseNet Output', use_column_width=True) | |
st.text("Results for image: %s" % filename) | |
st.text("Size of draw_image: {}".format(draw_image.shape)) | |
for pi in range(len(pose_scores)): | |
if pose_scores[pi] == 0.: | |
break | |
st.text('Pose #%d, score = %f' % (pi, pose_scores[pi])) | |
for ki, (s, c) in enumerate(zip(keypoint_scores[pi, :], keypoint_coords[pi, :, :])): | |
st.text('Keypoint %s, score = %f, coord = %s' % (posenet.PART_NAMES[ki], s, c)) | |
if __name__ == "__main__": | |
main() | |