Making sht_to_json compatible with the program

sht_to_json.py

@@ -1,23 +1,15 @@
 import sys
 import json
 
-filename = ""
-frame = {}
-shot = {}
-shots = []
-time = ""
-last_time = ""
-started = False
-
-def parse_shot_params(line):
+def parse_shot_params(line,shot):
     # assignments and list of frames
     for i in range(4, len(line)):
         if '=' in line[i]:
             sub_line = line[i].split('=')
             shot[sub_line[0]] = sub_line[1]
     shot['Frames'] = []
 
-def parse_map_of_values(line):
+def parse_map_of_values(line,frame):
     # map of values
     tag = line[3]
     dict = {}
@@ -27,32 +19,32 @@ def parse_map_of_values(line):
             dict[sub_line[i]] = float(sub_line[i + 1])
     frame[tag] = dict
 
-def parse_finetuned_shot_class(line):
+def parse_finetuned_shot_class(line,frame):
     # map of values
-    parse_map_of_values(line)
+    parse_map_of_values(line,frame)
 
-def parse_svm_shot_class(line):
+def parse_svm_shot_class(line,frame):
     # single string
     tag = line[3]
     frame[tag] = line[4]
 
-def parse_obj_class(line):
+def parse_obj_class(line,frame):
     # map of values
-    parse_map_of_values(line)
+    parse_map_of_values(line,frame)
 
-def parse_scene_location(line):
+def parse_scene_location(line,frame):
     # map of values
-    parse_map_of_values(line)
+    parse_map_of_values(line,frame)
 
-def parse_scene_attributes(line):
+def parse_scene_attributes(line,frame):
     # list of values
     tag = line[3]
     frame[tag] = []
     sub_line = line[4].split(', ')
     for i in range(0, len(sub_line)):
         frame[tag].append(sub_line[i])
 
-def parse_yolo_persons(line):
+def parse_yolo_persons(line,frame):
     if len(line) <= 5:
         return
     persons = {}
@@ -69,12 +61,14 @@ def parse_yolo_persons(line):
     frame['YOLO/PERSONS'] = persons
     pass
 
-if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print("Please write path to input file")
-        exit(1)
-    filename = sys.argv[1]
-
+def sht_to_json(sht_file_name):
+    frame = {}
+    shot = {}
+    shots = []
+    time = ""
+    last_time = ""
+    started = False
+    filename = sht_file_name
     for line in open(filename, 'r'):
         line = line.replace('\n', '').split('|')
         if len(line) > 3:
@@ -86,7 +80,7 @@ def parse_yolo_persons(line):
                 shot['Frames'].append(frame)
                 shots.append(shot)
             shot = {}
-            parse_shot_params(line)
+            parse_shot_params(line,shot)
             started = True
         elif started and len(line) > 3:
             # new frame
@@ -95,23 +89,23 @@ def parse_yolo_persons(line):
                     shot['Frames'].append(frame)
                     frame = {}
             tag = line[3]
-            if tag == 'FINETUNED_SHOT_CLASS':
-                parse_finetuned_shot_class(line)
-            elif tag == 'SVM_SHOT_CLASS':
-                parse_svm_shot_class(line)
+            if tag == 'FINETUNED_CLASS':
+                parse_finetuned_shot_class(line,frame)
+            elif tag == 'SVM_CLASS':
+                parse_svm_shot_class(line,frame)
             elif tag == 'OBJ_CLASS':
-                parse_obj_class(line)
+                parse_obj_class(line,frame)
             elif tag == 'SCENE_LOCATION':
-                parse_scene_location(line)
+                parse_scene_location(line,frame)
             elif tag == 'SCENE_ATTRIBUTES':
-                parse_scene_attributes(line)
+                parse_scene_attributes(line,frame)
             elif tag == 'YOLO/PERSONS':
-                parse_yolo_persons(line)
+                parse_yolo_persons(line,frame)
             else:
                 print("Unknown tag:" + tag)
                 exit(1)
 
-    with open(filename[:filename.index('.')] + '.json', 'w') as outfile:
+    with open(filename.split('.sht')[0] + '.json', 'w') as outfile:
         outfile.write("[\n")
         for i in range(0, len(shots)):
             json.dump(shots[i], outfile)