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() | |