forked from michael-mueller-git/Python-Funscript-Editor
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfunscript.py
More file actions
405 lines (296 loc) · 12.3 KB
/
Copy pathfunscript.py
File metadata and controls
405 lines (296 loc) · 12.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
""" Funscript Dataclass
Funscript is Funjack's haptic script format. It's basically JSON encoded timed positions:
.. code-block:: python
{
"version": "1.0",
"inverted": false,
"range": 90,
"actions": [
{"pos": 0, "at": 100},
{"pos": 100, "at": 500},
...
]
}
- version: funscript version (default="1.0")
- inverted (bool): positions are inverted, example: 0=100,100=0 (default=false)
- range: range of moment to use in percent (0-100) (default=90)
- actions: script for a Launch
- pos: position in percent (0-100)
- at : time to be at position in milliseconds
"""
import os
import cv2
import json
import numpy as np
class Funscript:
"""
Funscript Class
Args:
fps (float): Video FPS
version (str): funscript version (default="1.0")
inverted (bool): positions are inverted, example: 0=100,100=0 (default=false)
moment_range (int): range of moment to use in percent (0-100) (default=90)
"""
def __init__(self, fps, version='1.0', inverted=False, moment_range=90):
self.data = {
'version': version,
'inverted': inverted,
'range': max((0, min((moment_range, 100)))),
'fps': fps,
'actions': [],
'sections': []
}
self.changed = False
def is_inverted(self):
""" Check if script is inverted """
return self.data["inverted"]
def get_fps(self) -> float:
""" Get Video FPS
Returns:
float: video FPS
"""
return self.data['fps']
def is_empty(self) -> bool:
""" Check if the funscript actions are empty
Returns:
bool: True if no action exist else False
"""
return len(self.data['actions']) < 1
def delete_action(self, timestamp: int):
""" Deleta an action by timestamp
Args:
timestamp (int): timestamp in milliseconds
Returns:
Funscript: current funscript instance
"""
del_list = []
for i in range(len(self.data['actions'])):
if abs(self.data['actions'][i]['at'] - timestamp ) <= 2 * (1000.0 / self.data['fps']):
del_list.append(i)
del_list.sort(reverse=True)
for x in del_list: del self.data['actions'][x]
self.changed = True
return self
def delete_folowing_actions(self, timestamp: int):
""" Deleta all folowing actions for given timestamp
Args:
timestamp (int): timestamp in milliseconds
Returns:
Funscript: current funscript instance
"""
del_list = []
for i in range(len(self.data['actions'])):
if self.data['actions'][i]['at'] > timestamp+2:
del_list.append(i)
del_list.sort(reverse=True)
for x in del_list: del self.data['actions'][x]
self.changed = True
return self
def clear_actions(self):
""" Clear all actions
Returns:
Funscript: current funscript instance
"""
self.data['actions'] = []
return self
def get_last_action_time(self) -> int:
""" Get time of last action in current funscript
Returns:
int: timestamp in milliseconds
"""
if len(self.data['actions']) == 0: return 0
return self.data['actions'][-1]['at']
def get_first_action_time(self) -> int:
""" Get time of first action in current funscript
Returns:
int: timestamp in milliseconds
"""
if len(self.data['actions']) == 0: return 0
return self.data['actions'][0]['at']
def ground_all(self, limit :int = 45):
""" Set all lower strokes below the limit to zero
Args:
limit (int): all lower Strokes below this will be set to zero
Returns:
Funscript: current funscript instance
"""
self.changed = True
for i in range(1, len(self.data['actions'])-1):
if self.data['actions'][i]['pos'] < limit \
and self.data['actions'][i-1]['pos'] > self.data['actions'][i]['pos'] \
and self.data['actions'][i+1]['pos'] > self.data['actions'][i]['pos']: \
self.data['actions'][i]['pos'] = 0
return self
def get_next_action(self, current_timestamp: int) -> dict:
""" Get action next to current timestamp
Args:
current_timestamp (int): current timestamp in milliseconds
Returns:
dict: action dictionary with {'pos', 'at'}
"""
if len(self.data['actions']) < 1: return {'pos': 0, 'at': 0}
idx = (np.abs(np.array(self.get_actions_times()) - current_timestamp)).argmin()
if self.data['actions'][idx]['at'] > current_timestamp + 1: return self.data['actions'][idx]
elif len(self.data['actions']) == 1: return self.data['actions'][0]
elif idx+1 < len(self.data['actions']): return self.data['actions'][idx+1]
else: return self.data['actions'][0]
def get_prev_action(self, current_timestamp: int) -> dict:
""" Get previous action to current timestamp
Args:
current_timestamp (int): current timestamp in milliseconds
Returns:
dict: action dictionary with {'pos', 'at'}
"""
if len(self.data['actions']) < 2: return {'pos': 0, 'at': 0}
idx = (np.abs(np.array(self.get_actions_times()) - current_timestamp)).argmin()
if self.data['actions'][idx]['at'] < current_timestamp - 1: return self.data['actions'][idx]
elif idx > 0: return self.data['actions'][idx-1]
else: return self.data['actions'][-1]
def get_stroke_height(self, current_timestamp: int) -> int:
""" Get stroke height at given timestamp
Args:
current_timestamp (int): current timestamp in milliseconds
Returns:
int: stroke height (1-100)
"""
if len(self.get_actions()) < 2: return 0
return int(round(abs(self.get_next_action(current_timestamp)['pos'] - self.get_prev_action(current_timestamp)['pos'])))
def add_action(self, position: int, time: int, insert_raw: bool = False):
""" Add a new action to the Funscript
Args:
position (int): position in percent (0-100)
time (int): time to be at position in milliseconds
insert_raw (bool): do not delete near points
Returns:
Funscript: current funscript instance
"""
self.changed = True
if not insert_raw:
self.delete_action(time)
self.data['actions'].append({'pos': int(round(position)), 'at': time})
self.data['actions'].sort(key = lambda x: x['at'])
return self
def get_actions(self) -> list:
""" Get all actions from current funscript object
Returns:
list: funscript actions
"""
return self.data['actions']
def get_stroke_time(self, current_timestamp: int) -> int:
""" Get stroke duration for given timestamp in milliseconds
Note:
measure one periode (down-up-down or up-down-up)
Args:
current_timestamp (int): current position in milliseconds
Returns:
int: stroke duration in milliseconds
"""
stroke_times_before = [x for x in self.get_actions_times() if x <= current_timestamp]
stroke_times_after = [x for x in self.get_actions_times() if x > current_timestamp]
if len(stroke_times_before) == 0: stroke_times_before = [0]
if len(stroke_times_after) > 1: return int(round(stroke_times_after[1] - stroke_times_before[-1]))
elif len(stroke_times_before) > 2: return int(round(stroke_times_before[-1] - stroke_times_before[-3]))
else: return 0
def get_all_stroke_times(self) -> list:
""" Get all stroke duration in this Funscript
Returns:
list: list with stroke duration in milliseconds
"""
action_times = self.get_actions_times()
if len(action_times) < 2: return []
return [action_times[i+2] - action_times[i] for i in range(len(action_times)-2)]
def get_fastest_stroke(self) -> int:
""" Get the fastest stroke time in current Funscript
Returns:
int: fastest stroke time in milliseconds
"""
times = self.get_all_stroke_times()
return int(round(min(self.get_all_stroke_times()) if len(times) > 1 else 0))
def get_slowest_stroke(self) -> int:
""" Get the slowest stroke time in current Funscript
Returns:
int: slowest stroke time in milliseconds
"""
times = self.get_all_stroke_times()
return int(round(max(self.get_all_stroke_times()) if len(times) > 1 else 0))
def get_median_stroke(self) -> int:
""" Get the median stroke time for current Funscript
Returns:
int: median stroke time in milliseconds
"""
times = self.get_all_stroke_times()
return int(round(np.median(np.array(times)) if len(times) > 1 else 0))
def get_actions_positions(self) -> list:
""" Get all positions from current funscript object
Returns:
list: positions
"""
return [item['pos'] for item in self.get_actions()]
def get_actions_times(self) -> list:
""" Get all action times from current funscript object
Returns:
list: times in milliseconds
"""
return [item['at'] for item in self.get_actions()]
def __millisec_to_frame(self, milliseconds: int) -> int:
""" Convert milliseconds to frame number
Args:
milliseconds (int): time in milliseconds
Returns:
int: frame number for given time
"""
if milliseconds < 0: return 0
return int(round(float(milliseconds)/(float(1000)/float(self.data['fps']))))
def get_actions_frames(self) -> list:
""" Get all actions frame numbers from current funscript object
Returns:
list: frame numbers
"""
return [self.__millisec_to_frame(item['at']) for item in self.get_actions()]
def invert_actions(self):
""" Invert all actions in current Funscript
Returns:
Funscript: current funscript instance
"""
for i in range(len(self.data['actions'])):
self.data['actions'][i]['pos'] = round(100 - self.data['actions'][i]['pos'])
return self
def save(self, filename: str, create_backup: bool = True):
""" Save funscript to file
Args:
filename (path): path where to save the funscript
create_backup (bool): create an additional backup file
Returns:
Funscript: current funscript instance
"""
if not filename.endswith('.json') and not filename.endswith('.funscript'): filename += '.funscript'
for i in range(len(self.data['actions'])): self.data['actions'][i]['pos'] = round(self.data['actions'][i]['pos'])
with open(filename, 'w') as json_file: json.dump(self.data, json_file, indent=4)
if create_backup:
num=0
while os.path.exists(filename + str(num)): num += 1
filename += str(num)
with open(filename, 'w') as json_file: json.dump(self.data, json_file, indent=4) # save history
self.changed = False
return self
@staticmethod
def load(video_path: str, funscript_path: str):
""" Load funscript from file
Args:
funscript_path (path): funscript path
video_path (path): video path
Returns:
Funscript: a new funscript object
"""
with open(funscript_path, 'r') as json_file:
data = json.loads(json_file.read())
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
funscript = Funscript(fps=fps)
funscript.data = data
if 'inverted' in funscript.data.keys() and funscript.data['inverted'] == True:
funscript.invert_actions()
funscript.data['inverted'] = False
if 'fps' not in funscript.data.keys():
funscript.data['fps'] = fps
return funscript, funscript_path