| import gradio as gr |
| import cv2 |
| import tempfile |
| import numpy as np |
| import urllib.request |
| import os |
| import subprocess |
| import shutil |
| import mediapipe as mp |
| from mediapipe.tasks import python |
| from mediapipe.tasks.python import vision |
|
|
| model_url="https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/1/pose_landmarker_lite.task" |
| model_path="/tmp/pose_landmarker_lite.task" |
| if not os.path.exists(model_path): |
| urllib.request.urlretrieve(model_url,model_path) |
|
|
| class PoseDetector: |
| def __init__(self,confidence=0.25): |
| self.confidence=confidence |
| self.prev_landmarks=None |
| self.setup_detector(confidence) |
| def setup_detector(self,confidence): |
| base_options=python.BaseOptions(model_asset_path=model_path) |
| options=vision.PoseLandmarkerOptions( |
| base_options=base_options, |
| running_mode=vision.RunningMode.VIDEO, |
| num_poses=1, |
| min_pose_detection_confidence=confidence, |
| min_pose_presence_confidence=confidence, |
| min_tracking_confidence=confidence, |
| output_segmentation_masks=True) |
| self.pose_landmarker=vision.PoseLandmarker.create_from_options(options) |
| def enhance_image(self,frame): |
| lab=cv2.cvtColor(frame,cv2.COLOR_BGR2LAB) |
| l,a,b=cv2.split(lab) |
| clahe=cv2.createCLAHE(clipLimit=2.0,tileGridSize=(8,8)) |
| l=clahe.apply(l) |
| enhanced=cv2.merge([l,a,b]) |
| enhanced=cv2.cvtColor(enhanced,cv2.COLOR_LAB2BGR) |
| kernel=np.array([[-0.5,-0.5,-0.5],[-0.5,5,-0.5],[-0.5,-0.5,-0.5]]) |
| enhanced=cv2.filter2D(enhanced,-1,kernel) |
| enhanced=cv2.convertScaleAbs(enhanced,alpha=1.1,beta=10) |
| return enhanced |
| def remove_background(self,frame,segmentation_mask): |
| if segmentation_mask is None: |
| return frame |
| try: |
| mask=segmentation_mask.numpy_view() |
| if mask.shape[:2]!=(frame.shape[0],frame.shape[1]): |
| mask=cv2.resize(mask,(frame.shape[1],frame.shape[0])) |
| binary_mask=(mask>0.5).astype(np.uint8)*255 |
| kernel=np.ones((5,5),np.uint8) |
| binary_mask=cv2.morphologyEx(binary_mask,cv2.MORPH_CLOSE,kernel) |
| binary_mask=cv2.morphologyEx(binary_mask,cv2.MORPH_OPEN,kernel) |
| background=np.zeros_like(frame) |
| mask_3ch=cv2.cvtColor(binary_mask,cv2.COLOR_GRAY2BGR)/255.0 |
| result=(frame*mask_3ch+background*(1-mask_3ch)).astype(np.uint8) |
| return result |
| except: |
| return frame |
| def temporal_smoothing(self,current_landmarks,alpha=0.7): |
| if self.prev_landmarks is None or len(self.prev_landmarks)!=len(current_landmarks): |
| self.prev_landmarks=current_landmarks |
| return current_landmarks |
| smoothed=[] |
| for i in range(len(current_landmarks)): |
| class SmoothLandmark:pass |
| landmark=SmoothLandmark() |
| landmark.x=alpha*current_landmarks[i].x+(1-alpha)*self.prev_landmarks[i].x |
| landmark.y=alpha*current_landmarks[i].y+(1-alpha)*self.prev_landmarks[i].y |
| landmark.z=alpha*current_landmarks[i].z+(1-alpha)*self.prev_landmarks[i].z |
| landmark.visibility=current_landmarks[i].visibility |
| smoothed.append(landmark) |
| self.prev_landmarks=smoothed |
| return smoothed |
| def detect_pose_multi_pass(self,frame,timestamp_ms): |
| rgb_frame=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) |
| mp_image=mp.Image(image_format=mp.ImageFormat.SRGB,data=rgb_frame) |
| result=self.pose_landmarker.detect_for_video(mp_image,timestamp_ms) |
| if result.pose_landmarks: |
| return result,frame |
| enhanced_frame=self.enhance_image(frame) |
| rgb_enhanced=cv2.cvtColor(enhanced_frame,cv2.COLOR_BGR2RGB) |
| mp_image_enhanced=mp.Image(image_format=mp.ImageFormat.SRGB,data=rgb_enhanced) |
| result=self.pose_landmarker.detect_for_video(mp_image_enhanced,timestamp_ms) |
| if result.pose_landmarks: |
| return result,enhanced_frame |
| return None,frame |
|
|
| SIMPLER_CONNECTIONS=[(11,12),(11,23),(12,24),(23,24),(11,13),(13,15),(12,14),(14,16),(23,25),(25,27),(24,26),(26,28)] |
|
|
| def draw_landmarks(frame,landmarks): |
| h,w=frame.shape[:2] |
| for connection in SIMPLER_CONNECTIONS: |
| idx1,idx2=connection |
| if idx1<len(landmarks)and idx2<len(landmarks): |
| if landmarks[idx1].visibility>0.5 and landmarks[idx2].visibility>0.5: |
| pt1=(int(landmarks[idx1].x*w),int(landmarks[idx1].y*h)) |
| pt2=(int(landmarks[idx2].x*w),int(landmarks[idx2].y*h)) |
| if 0<=pt1[0]<w and 0<=pt1[1]<h and 0<=pt2[0]<w and 0<=pt2[1]<h: |
| cv2.line(frame,pt1,pt2,(0,255,0),3) |
| for i,landmark in enumerate(landmarks): |
| if landmark.visibility>0.5: |
| x=int(landmark.x*w) |
| y=int(landmark.y*h) |
| if 0<=x<w and 0<=y<h: |
| if i in[11,12]:color,radius=(0,165,255),6 |
| elif i in[23,24]:color,radius=(0,255,255),6 |
| elif i in[13,14,15,16]:color,radius=(255,0,0),5 |
| elif i in[25,26,27,28]:color,radius=(0,255,0),5 |
| else:color,radius=(0,0,255),4 |
| cv2.circle(frame,(x,y),radius,color,-1) |
| cv2.circle(frame,(x,y),radius,(255,255,255),1) |
|
|
| def process_video(input_video,remove_background=False,progress=gr.Progress()): |
| if input_video is None: |
| return None |
| temp_dir=tempfile.mkdtemp() |
| output_path=tempfile.NamedTemporaryFile(delete=False,suffix='.mp4') |
| output_path.close() |
| cap=cv2.VideoCapture(input_video) |
| if not cap.isOpened(): |
| raise ValueError("Could not open video") |
| fps=int(cap.get(cv2.CAP_PROP_FPS))or 30 |
| total_frames=int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| detector=PoseDetector(confidence=0.25) |
| frame_count=0 |
| poses_detected=0 |
| while True: |
| ret,frame=cap.read() |
| if not ret: |
| break |
| progress(frame_count/total_frames,desc=f"Frame {frame_count}/{total_frames}") |
| timestamp_ms=int(frame_count*(1000/fps)) |
| result,processed_frame=detector.detect_pose_multi_pass(frame,timestamp_ms) |
| if result and result.pose_landmarks: |
| poses_detected+=1 |
| for landmarks in result.pose_landmarks: |
| smoothed_landmarks=detector.temporal_smoothing(landmarks) |
| draw_landmarks(processed_frame,smoothed_landmarks) |
| if remove_background and result.segmentation_masks: |
| processed_frame=detector.remove_background(processed_frame,result.segmentation_masks[0]) |
| frame_path=os.path.join(temp_dir,f"frame_{frame_count:06d}.jpg") |
| cv2.imwrite(frame_path,processed_frame) |
| frame_count+=1 |
| cap.release() |
| progress(0.95,desc="Encoding video...") |
| cmd=['ffmpeg','-y','-framerate',str(fps),'-pattern_type','glob','-i',f'{temp_dir}/frame_*.jpg','-c:v','libx264','-profile:v','main','-level','3.0','-pix_fmt','yuv420p','-movflags','+faststart','-crf','23','-preset','medium',output_path.name] |
| try: |
| result=subprocess.run(cmd,capture_output=True,text=True,timeout=300) |
| if result.returncode==0 and os.path.exists(output_path.name): |
| shutil.rmtree(temp_dir) |
| return output_path.name |
| else: |
| return input_video |
| except: |
| shutil.rmtree(temp_dir) |
| return input_video |
|
|
| def get_default_video(): |
| video_list=['obj_01.mp4','obj_04.mp4','obj_05.mp4','obj_07.mp4'] |
| for video_path in video_list: |
| if os.path.exists(video_path): |
| return video_path |
| return None |
|
|
| with gr.Blocks(title="FormAI - Body Movement Tracker",theme=gr.themes.Soft())as demo: |
| gr.Markdown("# 🏋️ FormAI - Your Virtual Movement Coach\n\n**Choose tracking mode:**\n- **Standard Tracking**: Faster processing, keeps original background\n- **Background Removal**: Better for complex backgrounds, creates clean output\n") |
| with gr.Row(): |
| with gr.Column(scale=2): |
| video_input=gr.Video(label="Select Video",interactive=True,sources=["upload"],height=350,value=get_default_video()) |
| gr.Markdown("### Sample Videos") |
| with gr.Row(): |
| videos=[('01','obj_01.mp4'),('03','obj_04.mp4'),('05','obj_05.mp4'),('07','obj_07.mp4')] |
| for label,path in videos: |
| if os.path.exists(path): |
| btn=gr.Button(f"🎬 Video {label}",variant="secondary",size="sm") |
| btn.click(fn=lambda p=path:p,outputs=video_input) |
| gr.Markdown("### Tracking Options") |
| with gr.Row(): |
| track_btn=gr.Button("🎯 Standard Tracking",variant="primary",size="lg") |
| track_bg_btn=gr.Button("🎯 Track with Background Removal",variant="secondary",size="lg") |
| with gr.Column(scale=3): |
| video_output=gr.Video(label="Tracked Output",height=500,autoplay=True) |
| if get_default_video(): |
| demo.load(fn=lambda:process_video(get_default_video(),False),outputs=video_output) |
| track_btn.click(fn=lambda vid:process_video(vid,False),inputs=video_input,outputs=video_output) |
| track_bg_btn.click(fn=lambda vid:process_video(vid,True),inputs=video_input,outputs=video_output) |
|
|
| if __name__=="__main__": |
| demo.launch(server_name="0.0.0.0",server_port=7860) |
