app restoration

app.py

@@ -1,227 +1,154 @@
+"""Application to demo inpainting, Median and Bilateral Blur using streamlit.
+
+Run using: streamlit run 10_04_image_restoration_app.py
+"""
+
+import streamlit as st
+import pathlib
+from streamlit_drawable_canvas import st_canvas
 import cv2
-import mediapipe as mp
-import random
-import time
-import datetime  # Import datetime for timestamping video files
-
-# Initialize Mediapipe Pose
-mp_pose = mp.solutions.pose
-pose = mp_pose.Pose()
-mp_drawing = mp.solutions.drawing_utils
-
-# Initialize the video capture
-cap = cv2.VideoCapture(0)  # Use 0 for the default camera or provide a video file path
-
-# Set the resolution to 720p
-cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)  # Width for 720p
-cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 960)   # Height for a resolution higher than 720p
-
-# Define actions and their corresponding scores
-actions = ['Raise Right Hand', 'Raise Left Hand', 'Wave Right Hand', 'Wave Left Hand', 'Clap', 'Touch shoulder Right', 'Touch shoulder Left', 'Tilt Head Right', 'Tilt Head Left', 'Spread Hands', 'Spin'] # Updated action names
-# actions = ['Spin']#, 'Jump'] # Updated action names
-current_action = random.choice(actions)
-score = 0
-last_action_time = time.time()
-action_duration = 10  # seconds
-pause_duration = .5  # Define the pause duration
-
-# Add a boolean variable to check if pose estimation is enabled
-pose_estimation_enabled = True  # Set to True to enable pose estimation
-pose_estimation_paused_time = None  # Variable to track when pose estimation was paused
-
-# Initialize video writer
-video_writer = None  # Variable to hold the video writer
-is_recording = False  # Variable to track recording state
-
-# Add a variable to control pose drawing
-pose_drawing_enabled = False  # Add a variable to control pose drawing
-
-
-def ChangeAction():
-    global current_action, last_action_time, score, pose_estimation_enabled, pose_estimation_paused_time  # Declare global variables
-    new_action = random.choice(actions)  # Select a new action
-    # while new_action == current_action:  # Ensure the new action is different from the current one
-    #     new_action = random.choice(actions)
-    current_action = new_action  # Update current action
-    print(f"New action: {current_action}")  # Update text for new pose
-    last_action_time = time.time()
-    pose.reset()  # Reset pose values (assuming a reset method exists)
-    
-    pose_estimation_enabled = False  # Disable pose estimation
-    pose_estimation_paused_time = time.time()  # Update the pause time
-while True:
-    ret, frame = cap.read()
-    if not ret:
-        break
-
-    # Check if pose estimation was paused for more than pause_duration
-    if pose_estimation_paused_time and time.time() - pose_estimation_paused_time > pause_duration:
-        pose_estimation_enabled = True  # Re-enable pose estimation
-        pose_estimation_paused_time = None  # Reset the pause time
-
-    action_detected = False  # Replace with actual detection logic
-    # Process the frame with Mediapipe Pose only if enabled
-    if pose_estimation_enabled:
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        results = pose.process(frame_rgb)
-
-        # Draw pose landmarks if enabled
-        if pose_drawing_enabled and results.pose_landmarks:  # Check if drawing is enabled
-            mp_drawing.draw_landmarks(frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
-
-        if results.pose_landmarks:  # Check if pose landmarks are detected
-            # Get the coordinates of the landmarks
-            landmarks = results.pose_landmarks.landmark
-            
-            # Define required variables for actions
-            right_shoulder = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value]
-            right_wrist = landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value]
-            left_shoulder = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value]
-            left_wrist = landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value]
-            right_hip = landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value]
-            left_hip = landmarks[mp_pose.PoseLandmark.LEFT_HIP.value]
-            right_pinky = landmarks[mp_pose.PoseLandmark.RIGHT_PINKY.value]
-            right_index = landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value]
-            right_thumb = landmarks[mp_pose.PoseLandmark.RIGHT_THUMB.value]
-            left_pinky = landmarks[mp_pose.PoseLandmark.LEFT_PINKY.value]
-            left_index = landmarks[mp_pose.PoseLandmark.LEFT_INDEX.value]
-            left_thumb = landmarks[mp_pose.PoseLandmark.LEFT_THUMB.value]
-            nose = landmarks[mp_pose.PoseLandmark.NOSE.value]
-            near_shoulder_distance = 0.1  # Define the near shoulder distance
-            
-            # Check for each action
-            if current_action == 'Raise Right Hand' and right_wrist.y < nose.y:  # Condition for raising right hand
-                action_detected = True
-                print("Raise Right Hand")
-            elif current_action == 'Raise Left Hand' and left_wrist.y < nose.y:  # Condition for raising left hand
-                action_detected = True
-                print("Raise Left Hand")
-            elif current_action == 'Wave Right Hand' and (right_wrist.x > right_shoulder.x and right_wrist.y < nose.y):  # Condition for waving, adjusted to be further from shoulder
-                action_detected = True
-                print("Wave Right Hand")
-            elif current_action == 'Wave Left Hand' and (left_wrist.x < left_shoulder.x and left_wrist.y < nose.y):  # Condition for waving, adjusted to be further from shoulder
-                action_detected = True
-                print("Wave Left Hand")
-            elif current_action == 'Clap' and (abs(right_wrist.x - left_wrist.x) < 0.1 and right_wrist.y < left_wrist.y and left_wrist.y < right_shoulder.y):  # Condition for clapping based on proximity
-                action_detected = True
-                print("Clap")
-            elif current_action == 'Touch shoulder Right' and (
-                (((right_pinky.x - right_shoulder.x) ** 2 + (right_pinky.y - right_shoulder.y) ** 2) ** 0.5 < near_shoulder_distance) or 
-                (((right_index.x - right_shoulder.x) ** 2 + (right_index.y - right_shoulder.y) ** 2) ** 0.5 < near_shoulder_distance) or 
-                (((right_thumb.x - right_shoulder.x) ** 2 + (right_thumb.y - right_shoulder.y) ** 2) ** 0.5 < near_shoulder_distance)):  # Condition for touching shoulder
-                action_detected = True
-                print("Touch shoulder Right")
-            elif current_action == 'Touch shoulder Left' and (
-                (((left_pinky.x - left_shoulder.x) ** 2 + (left_pinky.y - left_shoulder.y) ** 2) ** 0.5 < near_shoulder_distance) or 
-                (((left_index.x - left_shoulder.x) ** 2 + (left_index.y - left_shoulder.y) ** 2) ** 0.5 < near_shoulder_distance) or 
-                (((left_thumb.x - left_shoulder.x) ** 2 + (left_thumb.y - left_shoulder.y) ** 2) ** 0.5 < near_shoulder_distance)):  # Condition for touching shoulder
-                action_detected = True
-                print("Touch shoulder Left")
-            elif current_action == 'Tilt Head Right':
-                is_nose_near_right_shoulder = (abs(nose.x - right_shoulder.x) < near_shoulder_distance) and (abs(nose.y - right_shoulder.y) < near_shoulder_distance)
-                is_nose_below_left_shoulder = nose.y > left_shoulder.y
-                
-                if is_nose_near_right_shoulder or is_nose_below_left_shoulder:
-                    action_detected = True
-                    print("Tilt Head Right")
-                    
-            elif current_action == 'Tilt Head Left':
-                is_nose_near_left_shoulder = (abs(nose.x - left_shoulder.x) < near_shoulder_distance) and (abs(nose.y - left_shoulder.y) < near_shoulder_distance)
-                is_nose_above_right_shoulder = nose.y > right_shoulder.y
-                
-                if is_nose_near_left_shoulder or is_nose_above_right_shoulder:
-                    action_detected = True
-                    print("Tilt Head Left")
-            elif current_action == 'Spread Hands':
-                is_right_hand_spread = right_wrist.x < right_shoulder.x and abs(right_wrist.y - right_shoulder.y) < near_shoulder_distance
-                is_left_hand_spread = left_wrist.x > left_shoulder.x and abs(left_wrist.y - left_shoulder.y) < near_shoulder_distance
-                
-                if is_right_hand_spread and is_left_hand_spread:
-                    action_detected = True
-                    print("Spread Hands")
-            elif current_action == 'Spin' and right_shoulder.x > left_shoulder.x:  # Condition for spinning
-                action_detected = True
-                print("Spin")
-            elif current_action == 'Jump' and (right_hip.y < left_hip.y):  # Condition for jumping
-                action_detected = True
-                print("Jump")
-
-        # Update action based on detection
-        if action_detected:
-            score += 1
-            text_to_display = f'{current_action} success!'
-            ChangeAction()  # Call the new function
+import numpy as np
+import io
+import base64
+from PIL import Image
+
+
+# Function to create a download link for output image
+def get_image_download_link(img, filename, text):
+    """Generates a link to download a particular image file."""
+    buffered = io.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
+
+
+# Set title.
+st.sidebar.title('Image Restoration')
+
+
+# Specify canvas parameters in application
+uploaded_file = st.sidebar.file_uploader("Upload Image to restore OK:", type=["png", "jpg"])
+image = None
+res = None
+
+if uploaded_file is not None:
+    # Debug: Print uploaded file information
+    # st.write("Uploaded file:", uploaded_file.name)
+
+    # Convert the file to an opencv image.
+    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
+    image = cv2.imdecode(file_bytes, 1)
+    # Debug: Print image shape
+    # st.write("Image shape:", image.shape)
+
+    # Display the uploaded image immediately
+    # st.image(image[:,:,::-1], caption='Uploaded Image')
+    # Display a selection box for choosing the filter to apply.
+    option = st.sidebar.selectbox('Median or Bilateral Blur or Inpaint?', ('None', 'Median Blur', 'Bilateral Blur', 'Image Inpaint'))
+
+    if option == 'Median Blur':
+        ksize = st.sidebar.slider("ksize: ", 3, 15, 5, 2)
+        image = cv2.medianBlur(image, ksize)
+        res=image[:,:,::-1]
+        st.image(res)
+    elif option == 'Bilateral Blur':
+        dia = st.sidebar.slider("diameter: ", 1, 50, 20)
+        sigmaColor = st.sidebar.slider("sigmaColor: ", 0, 250, 200, 10)
+        sigmaSpace = st.sidebar.slider("sigmaSpace: ", 0, 250, 100, 10)
+        image = cv2.bilateralFilter(image, dia, sigmaColor, sigmaSpace)
+        res=image[:,:,::-1]
+        st.image(res)
+
+    elif option == 'Image Inpaint':
+        # Debug: Print selected option
+        # st.write("Selected option for inpainting:", option)
+
+        stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
+        # st.write("Stroke width:", stroke_width)  # Debug: Print stroke width
+
+        h, w = image.shape[:2]
+        # st.write("Original image dimensions (h, w):", h, w)  # Debug: Print dimensions
+        if w > 800:
+            h_, w_ = int(h * 800 / w), 800
         else:
-            remaining_time = action_duration - (time.time() - last_action_time)
-            if remaining_time <= 0:
-                text_to_display = f'{current_action} missed!'
-                ChangeAction()  # Call the new function
-            else:
-                text_to_display = f'Time left to {current_action}: {int(remaining_time)}'
-
-    # # Get the size of the text to display the action result
-    # text_size = cv2.getTextSize(text_to_display, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0]  
-    # # Calculate the x position to center the text horizontally
-    # text_x = (frame.shape[1] - text_size[0]) // 2  
-    # Put the text on the frame at the bottom center
-    # cv2.putText(frame, text_to_display, 
-    #             (text_x, frame.shape[0] - 30),  # Centered at the bottom based on text width
-    #             cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)  # White color for visibility
-
-    # Get the size of the text for the current action to display
-    text_size = cv2.getTextSize(f'{current_action}', cv2.FONT_HERSHEY_SIMPLEX, 3, 5)[0]  
-    # Calculate the x position to center the current action text horizontally
-    text_x = (frame.shape[1] - text_size[0]) // 2  
-    # Put the current action text on the frame at the top center, adjusting the y position to avoid going off-screen
-    cv2.putText(frame, f'{current_action}', 
-                (text_x, text_size[1] + 10),  # Adjusted y position to ensure text is visible
-                cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 5, cv2.LINE_AA)  # Large font size and blue color for visibility
-    
-    # Create a banner for the score background
-    banner_height = 50
-    overlay = frame.copy()  # Create a copy of the frame for the overlay
-    cv2.rectangle(overlay, (0, frame.shape[0] - banner_height), (frame.shape[1], frame.shape[0]), (0, 0, 0), cv2.FILLED)  # Black banner at the bottom
-    alpha = 0.5  # Set transparency level
-    cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)  # Blend the overlay with the original frame
-
-    # Display the score at the bottom center of the frame
-    text_size = cv2.getTextSize(f'Score: {score}', cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0]  # Get the size of the text
-    text_x = (frame.shape[1] - text_size[0]) // 2  # Calculate the x position to center the text horizontally
-    cv2.putText(frame, f'Score: {score}', 
-                (text_x, frame.shape[0] - 20),  # Move text to the bottom center of the image
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)  # White text for visibility
-
-    # Show the frame
-    cv2.imshow('Body Tracking', frame)
-
-    key = cv2.waitKey(1) & 0xFF
-    # Check for quitting the application
-    if key == ord('q'):  # Press 'q' to quit
-        break
-    elif key == ord('r'):  # Press 'r' to reset score
-        score = 0
-    elif key == ord('v'):  # Press 'v' to toggle recording
-        if not is_recording:
-            user_input_name = input("Enter the name for the recording: ")  # Prompt user for input name
-            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")  # Get current timestamp
-            video_writer = cv2.VideoWriter(f'recording/{user_input_name}_{timestamp}.mp4', cv2.VideoWriter_fourcc(*'H264'), 20.0, (1280, 720))  # Initialize video writer in MP4 format with H.264 codec
-            is_recording = True
-            print("Recording started.")
+            h_, w_ = h, w
+        # st.write("Updated image dimensions (h_, w_):", h_, w_)  # Debug: Print dimensions
+
+        # Create a canvas component.
+        canvas_result = st_canvas(
+            fill_color='white',
+            stroke_width=stroke_width,
+            stroke_color='black',
+            background_image=Image.open(uploaded_file).resize((h_, w_)),
+            update_streamlit=True,
+            height=h_,
+            width=w_,
+            drawing_mode='freedraw',
+            key="canvas",
+        )
+        
+        # Debug: Print canvas result
+        # st.write("Canvas result:", canvas_result)
+
+        stroke = canvas_result.image_data
+        if stroke is not None:
+            # Debug: Print stroke data
+            # st.write("Stroke data shape:", stroke.shape)
+
+            if st.sidebar.checkbox('show mask'):
+                st.image(stroke)
+
+            mask = cv2.split(stroke)[3]
+            mask = np.uint8(mask)
+            mask = cv2.resize(mask, (w, h))
+            # Debug: Print mask shape
+            # st.write("Mask shape:", mask.shape)
+
+        st.sidebar.caption('Happy with the selection?')
+        option = st.sidebar.selectbox('Mode', ['None', 'Telea', 'NS', 'Compare both'])
+
+        if option == 'Telea':
+            st.subheader('Result of Telea')
+            res = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)[:,:,::-1]
+            st.image(res)
+            # Debug: Print result shape
+            # st.write("Telea result shape:", res.shape)
+        elif option == 'Compare both':
+            col1, col2 = st.columns(2)
+            res1 = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)[:,:,::-1]
+            res2 = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_NS)[:,:,::-1]
+            with col1:
+                st.subheader('Result of Telea')
+                st.image(res1)
+            with col2:
+                st.subheader('Result of NS')
+                st.image(res2)
+            if res1 is not None:
+                # Display link.
+                result1 = Image.fromarray(res1)
+                st.sidebar.markdown(
+                    get_image_download_link(result1, 'telea.png', 'Download Output of Telea'),
+                    unsafe_allow_html=True)
+            if res2 is not None:
+                # Display link.
+                result2 = Image.fromarray(res2)
+                st.sidebar.markdown(
+                    get_image_download_link(result2, 'ns.png', 'Download Output of NS'),
+                    unsafe_allow_html=True)
+
+        elif option == 'NS':
+            st.subheader('Result of NS')
+            res = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_NS)[:,:,::-1]
+            st.image(res)
         else:
-            is_recording = False
-            video_writer.release()  # Release the video writer
-            video_writer = None
-            print("Recording stopped.")
-    elif key == ord('d'):  # Press 'd' to toggle drawing of pose
-        pose_drawing_enabled = not pose_drawing_enabled  # Toggle the drawing state
-        print(f"Pose drawing {'enabled' if pose_drawing_enabled else 'disabled'}.")  # Print the current state
-
-    # Write the frame to the video file if recording
-    if is_recording and video_writer is not None:
-        video_writer.write(frame)
-
-pose.close()
-cap.release()
-if video_writer is not None:
-    video_writer.release()  # Ensure the video writer is released if still open
-cv2.destroyAllWindows()
+            pass
+
+    if res is not None:
+        # Debug: Print final result shape
+        # st.write("Final result shape:", res.shape)
+        # Display link.
+        result = Image.fromarray(res)
+        st.sidebar.markdown(
+            get_image_download_link(result, 'output.png', 'Download Output'),
+            unsafe_allow_html=True)

requirements.txt

@@ -1,6 +1,5 @@
-# opencv-python 
 numpy
 streamlit
 opencv-python-headless
 pillow
-mediapipe==0.10.14
+streamlit_drawable_canvas