Update app.py

app.py

@@ -2,21 +2,30 @@
 import os
 import json
 import numpy as np
+import matplotlib.pyplot as plt
 from pathlib import Path
 from pprint import pprint
 from omegaconf import OmegaConf
 from PIL import Image, ImageDraw
 import streamlit as st
+import random
 # %%
+os.environ['ROOT'] = os.path.dirname(os.path.realpath(__file__))#'/mnt/Documents/traffic_var_server/visualization'
+# print("os.environ['ROOT'] :",os.environ['ROOT'])
+# %%
+
+def get_list_folder(PATH):
+    return [name for name in os.listdir(PATH) if os.path.isdir(os.path.join(PATH, name))]
+
+def get_file_only(PATH):
+    return [f for f in os.listdir(PATH) if os.path.isfile(os.path.join(PATH, f))]
 
-os.environ['ROOT'] = os.path.dirname(os.path.realpath(__file__))
-print("os.environ['ROOT'] :",os.environ['ROOT'])
 # %%
 class ImageRetriever:
 
     def __init__(self, root_path, anno_path):
         self.root_path = Path(root_path)
-        self.anno = json.load(open(anno_path))
+        self.anno_path = Path(anno_path)
 
     def key2img_path(self, key):
         file_paths = [
@@ -31,21 +40,18 @@ class ImageRetriever:
             self.root_path / f"{key}.png",
             self.root_path / f"{key}.jpg",
         ]
-        print("file_paths!!!!!!!!", file_paths)
         for file_path in file_paths:
             if file_path.exists():
                 return file_path
 
 
-    def key2img(self, key, draw_bbox=True):
+    def key2img(self, key, temp_data, draw_bbox=True):
         file_path = self.key2img_path(key)
 
-        print("file_path!!@@@@", key, file_path)
-
         image = Image.open(file_path)
+
         if draw_bbox:
-            meta = self.anno[key]['details'][-1]
-            bboxes = [meta['bounding_box'].get(str(box_idx + 1), None) for box_idx in range(3)]
+            bboxes = [temp_data['bounding_box'].get(str(box_idx + 1), None) for box_idx in range(3)]
             image = self.hide_region(image, bboxes)
         return image
 
@@ -88,11 +94,9 @@ class ImageRetriever:
             image = Image.alpha_composite(image, overlay)
         return image
     
-def retrive_data(file_index, mode='direct'):
-    split = 'val'
-    mode = mode.lower()
-    main_dataset = json.load(open(os.path.join(os.environ['ROOT'], f"data/anno_{split}_{mode}.json")))
-    temp_data = main_dataset[list(main_dataset.keys())[file_index]]['details'][-1]
+def retrive_data(temp_data, img_key, mode='direct'):
+
+    # temp_data = main_dataset[list(main_dataset.keys())[file_index]]['details'][-1]
 
     message_dict = {}
 
@@ -108,29 +112,46 @@ def retrive_data(file_index, mode='direct'):
     message_dict['Entity #3'] = temp_data['Entity #3']
 
     img_retriever = ImageRetriever(
-    root_path=os.path.join(os.environ['ROOT'], ''),
-    anno_path=os.path.join(os.environ['ROOT'], f'data/anno_{split}_{mode}.json'),
+        root_path=os.path.join(os.environ['ROOT'], ''),
+        anno_path=os.path.join(os.environ['ROOT'], f'data/anno_{split}_{mode}.json'),
     )
-    img = img_retriever.key2img(list(main_dataset.keys())[file_index])
+    img = img_retriever.key2img(img_key, temp_data)
     img = img.resize((img.width // 2, img.height // 2))
 
     return img, message_dict
 
+# %%
+
+
 # %%
 if __name__ == '__main__':
     st.title("DHPR: Driving Hazard Prediction and Reasoning")
-    st.subheader("Data Visualization")
 
-    option = st.selectbox(
-    'Select the hazard type',
-    ('Direct', 'Indirect'))
+    img_path = os.path.join(os.environ['ROOT'], 'img/')
+    img_path_list = get_file_only(img_path)
+
+    split = 'val'
+    rand_index = 0
+    main_direct_dataset = json.load(open(os.path.join(os.environ['ROOT'], f"data/anno_{'val'}_{'direct'}.json")))
+    main_indirect_dataset = json.load(open(os.path.join(os.environ['ROOT'], f"data/anno_{'val'}_{'indirect'}.json")))
+
+    if st.button('Random Data Sample'):
+        rand_index = random.randint(0, len(get_file_only(img_path)))
+    else:
+        pass
+
+    st.subheader("Data Visualization")
 
-    st.write('You selected:', option)
+    img_key = img_path_list[rand_index].split('.')[0]
 
-    image_index = st.slider('Please Select The Image Index', 0, 999, 0)
-    st.write("You select the data index of ", image_index," for visualization from the validation set")
+    if img_key in main_direct_dataset.keys():
+        temp_data = main_direct_dataset[img_key]['details'][-1]
+    elif img_key in main_indirect_dataset.keys():
+        temp_data = main_indirect_dataset[img_key]['details'][-1]
+    else:
+        pass
 
-    img, message_dict =  retrive_data(image_index, option)
+    img, message_dict =  retrive_data(temp_data, img_key)
 
     st.write("---")