Commit
·
b668ddf
1
Parent(s):
93c1293
Add requirements.txt
Browse files- requirements.txt +5 -0
- utils.py +79 -0
requirements.txt
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch==2.2.1
|
2 |
+
ultralytics==8.0.186
|
3 |
+
streamlit
|
4 |
+
pillow==10.2.0
|
5 |
+
opencv-python-headless
|
utils.py
ADDED
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import numpy as np
|
3 |
+
import cv2
|
4 |
+
import tempfile, base64
|
5 |
+
import streamlit as st
|
6 |
+
import os
|
7 |
+
|
8 |
+
|
9 |
+
def readb64(uri):
|
10 |
+
encoded_data = uri.split(',')[-1]
|
11 |
+
nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8)
|
12 |
+
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
13 |
+
return img
|
14 |
+
|
15 |
+
def img2base64(img, extension="jpg"):
|
16 |
+
_, img_encoded = cv2.imencode(f".{extension}", img)
|
17 |
+
img_base64 = base64.b64encode(img_encoded)
|
18 |
+
img_base64 = img_base64.decode('utf-8')
|
19 |
+
return img_base64
|
20 |
+
|
21 |
+
def binary2video(video_binary):
|
22 |
+
temp_ = tempfile.NamedTemporaryFile(suffix='.mp4')
|
23 |
+
|
24 |
+
temp_.write(video_binary)
|
25 |
+
video_capture = cv2.VideoCapture(temp_.name)
|
26 |
+
ret, frame = video_capture.read()
|
27 |
+
return video_capture
|
28 |
+
|
29 |
+
def extract_frames(data_path, interval=30, max_frames=50):
|
30 |
+
"""Method to extract frames"""
|
31 |
+
cap = cv2.VideoCapture(data_path)
|
32 |
+
frame_num = 0
|
33 |
+
frames = list()
|
34 |
+
|
35 |
+
while cap.isOpened():
|
36 |
+
success, image = cap.read()
|
37 |
+
if not success:
|
38 |
+
break
|
39 |
+
if frame_num % interval == 0:
|
40 |
+
frames.append(image)
|
41 |
+
frame_num += 1
|
42 |
+
if len(frames) > max_frames:
|
43 |
+
break
|
44 |
+
cap.release()
|
45 |
+
return frames
|
46 |
+
|
47 |
+
def update_dir(key):
|
48 |
+
choice = st.session_state[key]
|
49 |
+
if os.path.isdir(os.path.join(st.session_state[key+'curr_dir'], choice)):
|
50 |
+
st.session_state[key+'curr_dir'] = os.path.normpath(os.path.join(st.session_state[key+'curr_dir'], choice))
|
51 |
+
files = sorted(os.listdir(st.session_state[key+'curr_dir']))
|
52 |
+
if "images" in files:
|
53 |
+
files.remove("images")
|
54 |
+
st.session_state[key+'files'] = files
|
55 |
+
|
56 |
+
def st_file_selector(st_placeholder, path='.', label='Select a file/folder', key = 'selected'):
|
57 |
+
if key+'curr_dir' not in st.session_state:
|
58 |
+
base_path = '.' if path is None or path == '' else path
|
59 |
+
base_path = base_path if os.path.isdir(base_path) else os.path.dirname(base_path)
|
60 |
+
base_path = '.' if base_path is None or base_path == '' else base_path
|
61 |
+
|
62 |
+
files = sorted(os.listdir(base_path))
|
63 |
+
files.insert(0, 'Choose a file...')
|
64 |
+
if "images" in files:
|
65 |
+
files.remove("images")
|
66 |
+
st.session_state[key+'files'] = files
|
67 |
+
st.session_state[key+'curr_dir'] = base_path
|
68 |
+
else:
|
69 |
+
base_path = st.session_state[key+'curr_dir']
|
70 |
+
|
71 |
+
selected_file = st_placeholder.selectbox(label=label,
|
72 |
+
options=st.session_state[key+'files'],
|
73 |
+
key=key,
|
74 |
+
on_change = lambda: update_dir(key))
|
75 |
+
|
76 |
+
if selected_file == "Choose a file...":
|
77 |
+
return None
|
78 |
+
|
79 |
+
return selected_file
|