Juartaurus/poi_Engineering
0
1import numpy as np2import gradio as gr3import cv24import os5import shutil6import re7import torch 8import csv9import time10 11from src.sts.demo.sts import handle_sts12from src.ir.ir import handle_ir13from src.ir.src.models.tc_classifier import TCClassifier14from src.tracker.signboard_track import SignboardTracker15 16from omegaconf import DictConfig17from hydra import compose, initialize18 19signboardTracker = SignboardTracker()20 21tracking_result_dir = ""22output_track_format = "mp4v"23output_track = ""24output_sts = ""25video_dir = ""26vd_dir = ""27labeling_dir = ""28 29frame_out = {}30rs = {}31results = []32 33# with initialize(version_base=None, config_path="src/ir/configs", job_name="ir"):34# config = compose(config_name="test")35# config: DictConfig36# model_ir = TCClassifier(config.model.train.model_name,37# config.model.train.n_classes,38# config.model.train.lr,39# config.model.train.scheduler_type,40# config.model.train.max_steps,41# config.model.train.weight_decay,42# config.model.train.classifier_dropout,43# config.model.train.mixout,44# config.model.train.freeze_encoder)45# model_ir = model_ir.load_from_checkpoint(checkpoint_path=config.ckpt_path, map_location=torch.device("cuda"))46 47def create_dir(list_dir_path):48 for dir_path in list_dir_path:49 if not os.path.isdir(dir_path):50 os.makedirs(dir_path)51 52def get_meta_from_video(input_video):53 if input_video is not None:54 video_name = os.path.basename(input_video).split('.')[0]55 56 global video_dir57 video_dir = os.path.join("static/videos/", f"{video_name}")58 59 global vd_dir60 vd_dir = os.path.join(video_dir, os.path.basename(input_video))61 62 global output_track63 output_track = os.path.join(video_dir,"original")64 65 global tracking_result_dir66 tracking_result_dir = os.path.join(video_dir,"track/cropped")67 68 global output_sts69 output_sts = os.path.join(video_dir,"track/sts")70 71 global labeling_dir72 labeling_dir = os.path.join(video_dir,"track/labeling")73 74 if os.path.isdir(video_dir):75 return None76 else:77 create_dir([output_track, video_dir, os.path.join(video_dir, "track/segment"), output_sts, tracking_result_dir, labeling_dir])78 79 # initialize the video stream80 video_cap = cv2.VideoCapture(input_video)81 # grab the width, height, and fps of the frames in the video stream.82 frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))83 frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))84 fps = int(video_cap.get(cv2.CAP_PROP_FPS))85 86 #tổng Fps87 # total_frames = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))88 # print(total_frames)89 # # Tính tổng số giây trong video90 # total_seconds = total_frames / video_cap.get(cv2.CAP_PROP_FPS)91 # print(total_seconds)92 93 # initialize the FourCC and a video writer object94 fourcc = cv2.VideoWriter_fourcc(*"mp4v")95 output = cv2.VideoWriter(vd_dir, fourcc, fps, (frame_width, frame_height))96 97 while True:98 success, frame = video_cap.read()99 # write the frame to the output file100 if success == True:101 output.write(frame)102 else:103 break104 # print(fps)105 # return gr.Slider(1, fps, value=4, label="FPS",step=1, info="Choose between 1 and {fps}", interactive=True)106 return gr.Textbox(value=fps)107 108def get_signboard(evt: gr.SelectData):109 name_fr = int(evt.index) + 1110 ids_dir = tracking_result_dir111 all_ids = os.listdir(ids_dir)112 gallery=[]113 for i in all_ids:114 fr_id = str(name_fr)115 al = re.search("[\d]*_"+fr_id+".png", i)116 if al:117 id_dir = os.path.join(ids_dir, i)118 gallery.append(id_dir)119 gallery = sorted(gallery)120 return gallery, name_fr121 122def tracking(fps_target):123 start = time.time()124 fps_target = int(fps_target)125 global results126 results = signboardTracker.inference_signboard(fps_target, vd_dir, output_track, output_track_format, tracking_result_dir)[0]127 # print("result", results)128 fd = []129 global frame_out130 list_id = []131 132 with open(os.path.join(video_dir, "track/label.csv"), 'w', newline='') as file:133 writer = csv.writer(file)134 writer.writerow(["Signboard", "Frame", "Text"])135 136 for frame, values in results.items():137 frame_dir = os.path.join(output_track, f"{frame}.jpg")138 # segment = os.path.join(video_dir,"segment/" + f"{frame}.jpg")139 list_boxs = []140 full = []141 list_id_tmp = []142 # print("values", values)143 for value in values:144 145 list_boxs.append(value['box']) 146 list_id_tmp.append(value['id'])147 _, dict_rec_sign_out = handle_sts(frame_dir, labeling_dir, list_boxs, list_id_tmp)148 149 # predicted = handle_ir(frame_dir, dict_rec_sign_out, os.path.join(video_dir, "ir"))150 # print(predicted)151 152 # fd.append(frame_dir)153 # frame_out[frame] = full154 list_id.extend(list_id_tmp)155 list_id = list(set(list_id))156 # print(list_id)157 print(time.time()-start)158 return gr.Dropdown(label="signboard",choices=list_id, interactive=True)159 160 161def get_select_index(img_id, evt: gr.SelectData):162 ids_dir = tracking_result_dir163 # print(ids_dir)164 all_ids = os.listdir(ids_dir)165 gallery = []166 for i in all_ids:167 fr_id = str(img_id)168 al = re.search("[\d]*_"+fr_id+".png", i)169 if al:170 id_dir = os.path.join(ids_dir, i)171 gallery.append(id_dir)172 gallery = sorted(gallery)173 174 gallery_id=[]175 id_name = gallery[evt.index]176 id = os.path.basename(id_name).split(".")[0].split("_")[0]177 for i in all_ids:178 al = re.search("^" +id + "_[\d]*.png", i)179 if al:180 id_dir = os.path.join(ids_dir, i)181 gallery_id.append(id_dir)182 gallery_id = sorted(gallery_id)183 return gallery_id184 185id_glb = None186def select_id(evt: gr.SelectData):187 choice=[]188 global id_glb189 id_glb = evt.value190 for key, values in results.items():191 for value in values:192 if value['id'] == evt.value:193 choice.append(int(key))194 return gr.Dropdown(label="frame", choices=choice, interactive=True)195 196 197import pandas as pd198 199frame_glb = None200def select_frame(evt: gr.SelectData):201 full_img = os.path.join(output_track, str(evt.value) + ".jpg")202 crop_img = os.path.join(tracking_result_dir, str(id_glb) + "_" + str(evt.value) + ".png")203 204 global frame_glb205 frame_glb = evt.value206 data = pd.read_csv(os.path.join(labeling_dir, str(id_glb) + "_" + str(frame_glb) + '.csv'), header=0)207 208 return full_img, crop_img, data209 210def get_data(dtfr):211 print(dtfr)212 213 # df = pd.read_csv(os.path.join(video_dir, "track/label.csv"))214 # for i, row in df.iterrows():215 # if str(row["Signboard"]) == str(id_tmp) and str(row["Frame"]) == str(frame_tmp):216 # # print(row["Text"])217 # df_new = df.replace(str(row["Text"]), str(labeling))218 # print(df_new)219 dtfr.to_csv(os.path.join(labeling_dir, str(id_glb) + "_" + str(frame_glb) + '.csv'), index=False, header=True)220 return221 222def seg_track_app():223 ##########################################################224 ###################### Front-end ########################225 ##########################################################226 with gr.Blocks(css=".gradio-container {background-color: white}") as demo:227 gr.Markdown(228 '''229 <div style="text-align:center;">230 <span style="font-size:3em; font-weight:bold;">POI Engineeing</span>231 </div>232 '''233 )234 with gr.Row():235 # video input236 with gr.Column(scale=0.2):237 238 tab_video_input = gr.Row(label="Video type input")239 with tab_video_input:240 input_video = gr.Video(label='Input video')241 242 tab_everything = gr.Row(label="Tracking")243 with tab_everything:244 with gr.Row():245 seg_signboard = gr.Button(value="Tracking", interactive=True)246 all_info = gr.Row(label="Information about video")247 with all_info:248 with gr.Row():249 text = gr.Textbox(label="Fps")250 check_fps = gr.Textbox(label="Choose fps for output", interactive=True)251 252 with gr.Column(scale=1):253 with gr.Row():254 with gr.Column(scale=2):255 with gr.Row():256 with gr.Column(scale=1):257 id_drop = gr.Dropdown(label="Signboards",choices=[])258 with gr.Column(scale=1):259 fr_drop = gr.Dropdown(label="Frames",choices=[])260 full_img = gr.Image(label="Full Image")261 262 with gr.Column(scale=1):263 crop_img = gr.Image(label="Cropped Image")264 with gr.Row():265 dtfr = gr.Dataframe(headers=["Tag", "Value"], datatype=["str", "str"], interactive=True)266 with gr.Row():267 submit = gr.Button(value="Submit", interactive=True)268 269 ##########################################################270 ###################### back-end #########################271 ##########################################################272 input_video.change(273 fn=get_meta_from_video,274 inputs=input_video,275 outputs=text276 )277 seg_signboard.click(278 fn=tracking,279 inputs=check_fps,280 outputs=id_drop281 )282 283 id_drop.select(select_id, None, fr_drop)284 fr_drop.select(select_frame, None, [full_img,crop_img, dtfr])285 submit.click(get_data, dtfr, None)286 287 demo.queue(concurrency_count=1)288 demo.launch(debug=True, enable_queue=True, share=True)289 290if __name__ == "__main__":291 seg_track_app()292