update app.py #4

app.py

@@ -22,6 +22,67 @@ model.load_state_dict(torch.load("model.pt"))
 model.eval()
 avg_pool = torch.nn.AdaptiveAvgPool2d((1, 1))
 
+import urllib.request
+urllib.request.urlretrieve("https://media.githubusercontent.com/media/commaai/speedchallenge/master/data/test.mp4", 'video_name.mp4')
+input_file = 'video_name.mp4'
+
+downsample_size = int(100)
+
+base_directory = os.getcwd()
+selected_directory = os.path.join(base_directory, "selected_images")
+if os.path.isdir(selected_directory):
+    shutil.rmtree(selected_directory)
+os.mkdir(selected_directory)
+
+zip_path = os.path.join(selected_directory, input_file.split('/')[-1][:-4] + ".zip")
+
+mean = np.asarray([0.3156024, 0.33569682, 0.34337464], dtype=np.float32)
+std = np.asarray([0.16568947, 0.17827448, 0.18925823], dtype=np.float32)
+
+img_vecs = []
+with torch.no_grad():
+    for fp_i, file_path in enumerate([input_file]):
+        for i, in_img in enumerate(video_reader(file_path,
+                                                targetFPS=9,
+                                                targetWidth=100,
+                                                to_rgb=True)):
+            in_img = (in_img.astype(np.float32) / 255.)
+            in_img = (in_img - mean) / std
+            in_img = np.expand_dims(in_img, 0)
+            in_img = np.transpose(in_img, (0, 3, 1, 2))
+            in_img = torch.from_numpy(in_img).float()
+            encoded = avg_pool(model(in_img))[0, :, 0, 0].cpu().numpy()
+            img_vecs += [encoded]
+img_vecs = np.asarray(img_vecs)
+print("images encoded")
+rv_indices, _ = furthest_neighbours(
+    img_vecs,
+    downsample_size,
+    seed=0)
+indices = np.zeros((img_vecs.shape[0],))
+indices[np.asarray(rv_indices)] = 1
+print("images selected")
+
+global_ctr = 0
+for fp_i, file_path in enumerate([input_file]):
+    for i, img in enumerate(video_reader(file_path,
+                                         targetFPS=9,
+                                         targetWidth=None,
+                                         to_rgb=False)):
+        if indices[global_ctr] == 1:
+            cv2.imwrite(join(selected_directory, str(global_ctr) + ".jpg"), img)
+        global_ctr += 1
+print("selected images extracted")
+
+all_selected_imgs_path = [join(selected_directory, f) for f in os.listdir(selected_directory) if
+                          isfile(join(selected_directory, f))]
+
+zipf = zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED)
+for i, f in enumerate(all_selected_imgs_path):
+    zipf.write(f, basename(f))
+zipf.close()
+print("selected images zipped")
+
 
 def predict(input_file, downsample_size):
     downsample_size = int(downsample_size)
@@ -32,7 +93,7 @@ def predict(input_file, downsample_size):
         shutil.rmtree(selected_directory)
     os.mkdir(selected_directory)
 
-    zip_path = os.path.join(input_file.split('/')[-1][:-4] + ".zip")
+    zip_path = os.path.join(selected_directory, input_file.split('/')[-1][:-4] + ".zip")
 
     mean = np.asarray([0.3156024, 0.33569682, 0.34337464], dtype=np.float32)
     std = np.asarray([0.16568947, 0.17827448, 0.18925823], dtype=np.float32)
@@ -90,4 +151,4 @@ demo = gr.Interface(
             gr.components.Number(label="Downsample size")],
     outputs=gr.components.File(label="Zip"))
 
-demo.launch(debug = True)
+demo.launch(debug=True)