max-unfinity commited on
Commit
a7a43e9
·
1 Parent(s): 8a7f63f
Files changed (1) hide show
  1. app.py +9 -10
app.py CHANGED
@@ -66,25 +66,24 @@ st.title('Pose Estimation App')
66
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
67
  st.caption(f'Using device: {device}')
68
 
69
- upload_tab, camera_tab = st.tabs(["Upload Photo", "Webcam Capture"])
70
 
71
- with upload_tab:
72
- uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
73
-
74
- with camera_tab:
75
- img_file_buffer = st.camera_input("Take a picture")
76
 
77
  img = None
78
- if img_file_buffer is not None:
79
- img = Image.open(img_file_buffer)
80
- if uploaded_file is not None:
81
- img = Image.open(uploaded_file)
82
 
 
83
  if img is not None:
84
  # predict
85
  with st.spinner('Predicting...'):
86
  model = load_model()
87
  pred = model(img)[0]
 
88
  keypoints = infer.get_keypoints(pred)
89
  if keypoints is not None:
90
  img = draw_output(img, keypoints)
 
66
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
67
  st.caption(f'Using device: {device}')
68
 
69
+ mode = st.radio('Select mode:', ['Upload an Image', 'Webcam Capture'])
70
 
71
+ if mode == 'Upload an Image':
72
+ img_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
73
+ elif mode == 'Webcam Capture':
74
+ img_file = st.camera_input("Take a picture")
 
75
 
76
  img = None
77
+ if img_file is not None:
78
+ img = Image.open(img_file)
 
 
79
 
80
+ st.divider()
81
  if img is not None:
82
  # predict
83
  with st.spinner('Predicting...'):
84
  model = load_model()
85
  pred = model(img)[0]
86
+ st.markdown('**Results:**')
87
  keypoints = infer.get_keypoints(pred)
88
  if keypoints is not None:
89
  img = draw_output(img, keypoints)