Update app.py

app.py

@@ -1,16 +1,261 @@
-import gradio as gr
+#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
+
+@st.cache_data()
+
+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 video_identity(video):
-    return video
+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))
 
-demo = gr.Interface(video_identity, 
-                    gr.Video(), 
-                    "playable_video",
-                    cache_examples=True)
+            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__":
-    demo.launch()
+    main()