Pratyush101's picture
Update app.py
470932b verified
raw
history blame
14.6 kB
import logging
import queue
from pathlib import Path
from typing import List, NamedTuple
import mediapipe as mp
import av
import cv2
import numpy as np
import streamlit as st
from streamlit_webrtc import WebRtcMode, webrtc_streamer
from sample_utils.download import download_file
from sample_utils.turn import get_ice_servers
# Logging setup
logger = logging.getLogger(__name__)
# Streamlit setup
st.title("AI Squat Detection using WebRTC")
st.info("Use your webcam for real-time squat detection.")
# Initialize MediaPipe components
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
class Detection(NamedTuple):
class_id: int
label: str
score: float
box: np.ndarray
# Angle calculation function
def calculate_angle(a, b, c):
a = np.array(a)
b = np.array(b)
c = np.array(c)
radians = np.arctan2(c[1]-b[1], c[0]-b[0]) - np.arctan2(a[1]-b[1], a[0]-b[0])
angle = np.abs(radians * 180.0 / np.pi)
if angle > 180.0:
angle = 360 - angle
return angle
# counterL=0#Counter checks for number of curls
# correct=0
# incorrect=0
# stage='mid'#it checks if we our hand is UP or DOWN
# Detection Queue
result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
counterL=0#Counter checks for number of curls
correct=0
incorrect=0
stage='mid'#it checks if we our hand is UP or DOWN
image = frame.to_ndarray(format="bgr24")
h, w = image.shape[:2]
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
results = pose.process(image_rgb)
landmarks = results.pose_landmarks.landmark if results.pose_landmarks else []
# Corrected detection logic
detections = [
Detection(
class_id=0, # Assuming a generic class_id for pose detections
label="Pose",
score=0.7, # Full confidence as pose landmarks were detected
box=np.array([0, 0, image.shape[1], image.shape[0]]) # Full image as bounding box
)
] if landmarks else []
if landmarks:
hipL = [landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y]
kneeL = [landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].y]
ankleL = [landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].y]
shoulderL = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
footIndexL = [landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value].y]
# Calculate angles
angleKneeL = calculate_angle(hipL, kneeL, ankleL)
angleHipL = calculate_angle(shoulderL, hipL, [hipL[0], 0])
angleAnkleL = calculate_angle(footIndexL, ankleL, kneeL)
#Visualize of left leg
cv2.putText(image, str(angleHipL),tuple(np.multiply(angleHipL, [640, 480]).astype(int)),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA)
# # Squat logic
# if 80 < angleKneeL < 110 and 29 < angleHipL < 40:
# cv2.putText(image, "Squat Detected!", (300, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 3)
# else:
# if angleHipL < 29:
# cv2.putText(image, "Lean Forward!", (300, 200), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
# elif angleHipL > 45:
# cv2.putText(image, "Lean Backward!", (300, 200), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
# if angleKneeL < 80:
# cv2.putText(image, "Squat Too Deep!", (300, 250), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
# elif angleKneeL > 110:
# cv2.putText(image, "Lower Your Hips!", (300, 300), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
# 1. Bend Forward Warning
if 10 < angleHipL < 18:
cv2.rectangle(image, (310, 180), (450, 220), (0, 0, 0), -1)
cv2.putText(image,f"Bend Forward",(320,200),cv2.FONT_HERSHEY_SIMPLEX,1,(150,120,255),1,cv2.LINE_AA)
# 2. Lean Backward Warning
if angleHipL > 45:
cv2.rectangle(image, (310, 180), (450, 220), (0, 0, 0), -1)
cv2.putText(image,f"Bend Backward",(320,200),cv2.FONT_HERSHEY_SIMPLEX,1,(80,120,255),1,cv2.LINE_AA)
# # stage 2
# # Incorrect movements
# 3. Knees not low enough
if 110 < angleKneeL < 130:
cv2.rectangle(image, (220, 40), (450, 80), (0, 0, 0), -1)
cv2.putText(image,f"Lower Your Hips",(230,60),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),1,cv2.LINE_AA)
# # 3. Knees not low enough and not completed the squat
# if angleKneeL>130 and stage=='mid':
# cv2.rectangle(image, (220, 40), (450, 80), (0, 0, 0), -1)
# cv2.putText(image,f"Lower Your Hips",(230,60),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),1,cv2.LINE_AA)
# incorrect+=1
# stage='up'
# # 4. Squat too deep
# if angleKneeL < 80 and stage=='mid':
# cv2.rectangle(image, (220, 40), (450, 80), (0, 0, 0), -1)
# cv2.putText(image,f"Squat too deep",(230,60),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),1,cv2.LINE_AA)
# incorrect +=1
# stage='up'
# stage 4
if (80 < angleKneeL < 110):
# if (18 < angleHipL < 40): # Valid "down" position
correct+=1
# stage='up'
# cv2.putText(image,f"Correct:{correct}",
# (400,120),cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,0),1,cv2.LINE_AA)
# cv2.putText(image,f"Incorrect:{incorrect}",
# (400,150),cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,0),1,cv2.LINE_AA)
#Render Counter to our camera screen
#Setup Status box
cv2.rectangle(image,(0,0),(500,80),(245,117,16),-1)
#REP data
cv2.putText(image,'Left',(10,12),
cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,0),1,cv2.LINE_AA)
cv2.putText(image,str(correct),
(10,60),cv2.FONT_HERSHEY_SIMPLEX,2,(255,255,255),2,cv2.LINE_AA)
#Stage data for left leg
cv2.putText(image,'STAGE',(230,12),
cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,0),1,cv2.LINE_AA)
cv2.putText(image,stage,
(230,60),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),1,cv2.LINE_AA)
mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,mp_drawing.DrawingSpec(color=(255, 175, 0), thickness=2, circle_radius=2),mp_drawing.DrawingSpec(color=(0, 255, 200), thickness=2, circle_radius=2))
result_queue.put(detections)
return av.VideoFrame.from_ndarray(image, format="bgr24")
# WebRTC streamer configuration
webrtc_streamer(
key="squat-detection",
mode=WebRtcMode.SENDRECV,
rtc_configuration={"iceServers": get_ice_servers(), "iceTransportPolicy": "relay"},
media_stream_constraints={"video": True, "audio": False},
video_frame_callback=video_frame_callback,
async_processing=True,
)
# import logging
# import cv2
# import numpy as np
# import streamlit as st
# from streamlit_webrtc import WebRtcMode, webrtc_streamer
# from cvzone.HandTrackingModule import HandDetector
# from cvzone.SelfiSegmentationModule import SelfiSegmentation
# import os
# import time
# import av
# import queue
# from typing import List, NamedTuple
# from sample_utils.turn import get_ice_servers
# logger = logging.getLogger(__name__)
# # Streamlit settings
# st.set_page_config(page_title="Virtual Keyboard", layout="wide")
# st.title("Interactive Virtual Keyboard")
# st.subheader('''Turn on the webcam and use hand gestures to interact with the virtual keyboard.
# Use 'a' and 'd' from the keyboard to change the background.''')
# # Initialize modules
# detector = HandDetector(maxHands=1, detectionCon=0.8)
# segmentor = SelfiSegmentation()
# # Define virtual keyboard layout
# keys = [["Q", "W", "E", "R", "T", "Y", "U", "I", "O", "P"],
# ["A", "S", "D", "F", "G", "H", "J", "K", "L", ";"],
# ["Z", "X", "C", "V", "B", "N", "M", ",", ".", "/"]]
# class Button:
# def __init__(self, pos, text, size=[100, 100]):
# self.pos = pos
# self.size = size
# self.text = text
# class Detection(NamedTuple):
# label: str
# score: float
# box: np.ndarray
# # result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
# listImg = os.listdir('model/street') if os.path.exists('model/street') else []
# if not listImg:
# st.error("Error: 'street' directory is missing or empty. Please add background images.")
# st.stop()
# else:
# imgList = [cv2.imread(f'model/street/{imgPath}') for imgPath in listImg if cv2.imread(f'model/street/{imgPath}') is not None]
# indexImg = 0
# prev_key_time = [time.time()] * 2
# output_text = ""
# if "output_text" not in st.session_state:
# st.session_state["output_text"] = ""
# # def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
# # img = frame.to_ndarray(format="bgr24")
# # hands, img = detector.findHands(img, flipType=False)
# # # Render hand detection results
# # if hands:
# # hand = hands[0]
# # bbox = hand["bbox"]
# # cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), (255, 0, 0), 2)
# # cv2.putText(img, 'OpenCV', (50,50), font,
# # fontScale, color, thickness, cv2.LINE_AA)
# # cv2.putText(img, 'OpenCV', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 1, cv2.LINE_AA)
# # result_queue.put(hands)
# # return av.VideoFrame.from_ndarray(img, format="bgr24")
# result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
# def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
# image = frame.to_ndarray(format="bgr24")
# # Run inference
# blob = cv2.dnn.blobFromImage(
# cv2.resize(image, (300, 300)), 0.007843, (300, 300), 127.5
# )
# net.setInput(blob)
# output = net.forward()
# h, w = image.shape[:2]
# # Convert the output array into a structured form.
# output = output.squeeze() # (1, 1, N, 7) -> (N, 7)
# output = output[output[:, 2] >= score_threshold]
# detections = [
# Detection(
# class_id=int(detection[1]),
# label=CLASSES[int(detection[1])],
# score=float(detection[2]),
# box=(detection[3:7] * np.array([w, h, w, h])),
# )
# for detection in output
# ]
# # Render bounding boxes and captions
# for detection in detections:
# caption = f"{detection.label}: {round(detection.score * 100, 2)}%"
# color = COLORS[detection.class_id]
# xmin, ymin, xmax, ymax = detection.box.astype("int")
# cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
# cv2.putText(
# image,
# caption,
# (xmin, ymin - 15 if ymin - 15 > 15 else ymin + 15),
# cv2.FONT_HERSHEY_SIMPLEX,
# 0.5,
# color,
# 2,
# )
# result_queue.put(detections)
# return av.VideoFrame.from_ndarray(image, format="bgr24")
# # def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
# # global indexImg, output_text
# # img = frame.to_ndarray(format="bgr24")
# # imgOut = segmentor.removeBG(img, imgList[indexImg])
# # hands, imgOut = detector.findHands(imgOut, flipType=False)
# # buttonList = [Button([30 + col * 105, 30 + row * 120], key) for row, line in enumerate(keys) for col, key in enumerate(line)]
# # detections = []
# # if hands:
# # for i, hand in enumerate(hands):
# # lmList = hand['lmList']
# # bbox = hand['bbox']
# # label = "Hand"
# # score = hand['score']
# # box = np.array([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]])
# # detections.append(Detection(label=label, score=score, box=box))
# # if lmList:
# # x4, y4 = lmList[4][0], lmList[4][1]
# # x8, y8 = lmList[8][0], lmList[8][1]
# # distance = np.sqrt((x8 - x4) ** 2 + (y8 - y4) ** 2)
# # click_threshold = 10
# # for button in buttonList:
# # x, y = button.pos
# # w, h = button.size
# # if x < x8 < x + w and y < y8 < y + h:
# # cv2.rectangle(imgOut, button.pos, (x + w, y + h), (0, 255, 160), -1)
# # cv2.putText(imgOut, button.text, (x + 20, y + 70), cv2.FONT_HERSHEY_PLAIN, 5, (255, 255, 255), 3)
# # if (distance / np.sqrt(bbox[2] ** 2 + bbox[3] ** 2)) * 100 < click_threshold:
# # if time.time() - prev_key_time[i] > 2:
# # prev_key_time[i] = time.time()
# # if button.text != 'BS' and button.text != 'SPACE':
# # output_text += button.text
# # elif button.text == 'BS':
# # output_text = output_text[:-1]
# # else:
# # output_text += ' '
# # result_queue.put(detections)
# # st.session_state["output_text"] = output_text
# # return av.VideoFrame.from_ndarray(imgOut, format="bgr24")
# webrtc_streamer(
# key="virtual-keyboard",
# mode=WebRtcMode.SENDRECV,
# rtc_configuration={"iceServers": get_ice_servers(), "iceTransportPolicy": "relay"},
# media_stream_constraints={"video": True, "audio": False},
# video_frame_callback=video_frame_callback,
# async_processing=True,
# )
# st.subheader("Output Text")
# st.text_area("Live Input:", value=st.session_state["output_text"], height=200)