adding multiprocessing version of image analyze

.gitignore

@@ -1,3 +1,6 @@
+# Added
+tmp/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

app.py

@@ -28,8 +28,11 @@ import numpy as np
 import gradio as gr
 import glob
 from inference import inference_frame
+from inference import inference_frame_par_ready
+from inference import process_frame
 import os
 import pathlib
+import multiprocessing as mp
 
 from time import time
 
@@ -37,7 +40,7 @@ def analyze_video(x, skip_frames = 5, frame_rate_out = 8):
     print(x)
 
     #Define path to saved images
-    path = '/tmp/test/'
+    path = 'tmp/test/'
     os.makedirs(path, exist_ok=True)
     
     # Define name of current video as number of videos in path
@@ -56,34 +59,128 @@ def analyze_video(x, skip_frames = 5, frame_rate_out = 8):
     counter = 0
     
     while(cap.isOpened()):
+        frames = []
         start = time()
-        ret, frame = cap.read()
+        for i in range(16):
+            start = time()
+            ret, frame = cap.read()
+            frames.append(frame)
+            if ret == False:
+                break
         print(f'read time: {time()-start}')
         
-        if ret==True:
-            if counter % skip_frames == 0:
-                name = os.path.join(path,f'{counter:05d}.png')
-                start = time()
-                frame = inference_frame(frame)
-                print(f'inference time: {time()-start}')
-                # write the flipped frame
-                start = time()
-                cv2.imwrite(name, frame)
-                print(f'write time: {time()-start}')
-            else: 
-                pass
-            print(counter)
-            counter +=1 
-        else:
-            break
+        #if ret==True:
+
+            #if counter % skip_frames == 0:
+        name = os.path.join(path,f'{counter:05d}.png')
+        # Get timing for inference
+        start = time()
+        frames = inference_frame(frames)
+        print(f'inference time: {time()-start}')
+        # write the flipped frame
+        
+        start = time()
+        for frame in frames:
+            name = os.path.join(path,f'{counter:05d}.png')
+            cv2.imwrite(name, frame)
+            counter +=1
+        print(f'write time: {time()-start}')
+        # else: 
+            
+        # print(counter)
+        # counter +=1
+        # else:
+        #     break
+    
+    # Release everything if job is finished
+    cap.release()
+
+    # Create video from predicted images
+    print(path)
+    os.system(f'''ffmpeg -framerate {frame_rate_out} -pattern_type glob -i '{path}/*.png'  -c:v libx264 -pix_fmt yuv420p {outname} -y''')
+    return outname
+
+
+def analyze_video_parallel(x, skip_frames = 5, 
+                           frame_rate_out = 8, batch_size = 16):
+    print(x)
+
+    #Define path to saved images
+    path = '/tmp/test/'
+    os.makedirs(path, exist_ok=True)
+    
+    # Define name of current video as number of videos in path
+    n_videos_in_path = len(os.listdir(path))
+    path = f'{path}{n_videos_in_path}'
+    os.makedirs(path, exist_ok=True)
+    
+    # Define name of output video
+    outname = f'{path}_processed.mp4'
+    
+    if os.path.exists(outname):
+        print('video already processed')
+        return outname
     
+    cap = cv2.VideoCapture(x)
+    counter = 0
+    pred_results_all = []
+    frames_all = []
+    while(cap.isOpened()):
+        frames = []
+        #start = time()
+        
+        while len(frames) < batch_size:
+            #start = time()
+            ret, frame = cap.read()
+            if ret == False:
+                break
+            elif counter % skip_frames == 0:
+                frames.append(frame)
+            counter += 1
+
+        #print(f'read time: {time()-start}')
+
+        frames_all.extend(frames)
+
+        # Get timing for inference
+        start = time()
+        print('len frames passed: ', len(frames))
+        
+        if len(frames) > 0:
+            pred_results = inference_frame_par_ready(frames)
+            print(f'inference time: {time()-start}')
+            pred_results_all.extend(pred_results)
+
+        # break while loop when return of the image reader is False
+        if ret == False:
+            break
+
+    print('exited prediction loop')
     # Release everything if job is finished
     cap.release()
+        
+    start = time()
+    pool = mp.Pool(mp.cpu_count()-2)
+    pool_out = pool.map(process_frame, 
+                        list(zip(pred_results_all, 
+                                    frames_all, 
+                                    [i for i in range(len(pred_results_all))])))
+    pool.close()
+    print(f'pool time: {time()-start}')
+    
+    start = time()
+    counter = 0
+    for pool_out_tmp in pool_out:
+        name = os.path.join(path,f'{counter:05d}.png')
+        cv2.imwrite(name, pool_out_tmp)
+        counter +=1
+    print(f'write time: {time()-start}')
 
     # Create video from predicted images
     print(path)
     os.system(f'''ffmpeg -framerate {frame_rate_out} -pattern_type glob -i '{path}/*.png'  -c:v libx264 -pix_fmt yuv420p {outname} -y''')
     return outname
+  
 
 def set_example_image(example: list) -> dict:
     return gr.Video.update(value=example[0])
@@ -120,7 +217,7 @@ with gr.Blocks(title='Shark Patrol',theme=gr.themes.Soft(),live=True,) as demo:
                                     samples=[[path.as_posix()]
                                              for path in paths if 'videos_side_by_side' not in str(path)])
 
-    video_button.click(analyze_video, inputs=video_input, outputs=video_output)
+    video_button.click(analyze_video_parallel, inputs=video_input, outputs=video_output)
 
     example_images.click(fn=set_example_image,
                          inputs=example_images,

inference.py

@@ -13,6 +13,10 @@ from mmdet.registry import VISUALIZERS
 
 from huggingface_hub import hf_hub_download
 from huggingface_hub import snapshot_download
+from time import time
+
+import concurrent.futures
+import threading
 
 
 classes= ['Beach',
@@ -76,7 +80,7 @@ classes= ['Beach',
 REPO_ID = "SharkSpace/maskformer_model"
 FILENAME = "mask2former"
 
-snapshot_download(repo_id=REPO_ID, token= os.environ.get('SHARK_MODEL'),local_dir='model/')
+# snapshot_download(repo_id=REPO_ID, token= os.environ.get('SHARK_MODEL'),local_dir='model/')
 
 
 
@@ -94,20 +98,75 @@ register_all_modules()
 # build the model from a config file and a checkpoint file
 model = init_detector(config_file, checkpoint_file, device='cuda:0')  # or device='cuda:0'
 model.dataset_meta['classes'] = classes
+print(model.cfg.visualizer)
 # init visualizer(run the block only once in jupyter notebook)
 visualizer = VISUALIZERS.build(model.cfg.visualizer)
+print(dir(visualizer))
 # the dataset_meta is loaded from the checkpoint and
 # then pass to the model in init_detector
 visualizer.dataset_meta = model.dataset_meta
+
 def inference_frame(image):
+    #import ipdb; ipdb.set_trace()
     result = inference_detector(model, image)
     # show the results
+    #import ipdb; ipdb.set_trace()
+    frames = []
+    cnt=0
+    
+    for res in result:
+        visualizer.add_datasample(
+        'result',
+        image[cnt],
+        data_sample=res.numpy(),
+        draw_gt = None,
+        show=False
+        )
+        frame = visualizer.get_image()
+        frames.append(frame)
+        cnt+=1
+
+    #frames = process_frames(result, image, visualizer)
+    end = time()
+    print("Time taken for drawing: ", end-start)
+    return frames
+
+def inference_frame_par_ready(image):
+    result = inference_detector(model, image)
+    return [result[i].numpy() for i in range(len(result))]
+
+def process_frame(in_tuple = (None, None, None)):
     visualizer.add_datasample(
-    'result',
-    image,
-    data_sample=result,
-    draw_gt = None,
-    show=False
-    )
-    frame = visualizer.get_image()
-    return frame
+        'result',
+        in_tuple[1], #image,
+        data_sample=in_tuple[0], #result
+        draw_gt = None,
+        show=False
+        )
+    
+    #frame = visualizer.get_image()
+    #print(in_tuple[2])
+    return visualizer.get_image()
+
+#def process_frame(frame):
+
+# def process_frames(result, image, visualizer):
+#     frames = []
+#     lock = threading.Lock()
+
+#     def process_data(cnt, res, img):
+#         visualizer.add_datasample('result', img, data_sample=res, draw_gt=None, show=False)
+#         frame = visualizer.get_image()
+#         with lock:
+#             frames.append(frame)
+
+#     threads = []
+#     for cnt, res in enumerate(result):
+#         t = threading.Thread(target=process_data, args=(cnt, res, image[cnt]))
+#         threads.append(t)
+#         t.start()
+
+#     for t in threads:
+#         t.join()
+
+#     return frames