Vfrae/YOLO-FishScale
0
1''' app.py2>>> 💥 This file comprises the main function to run the demo of the project.3'''4 5#------------------------------------------------------------------------------#6# IMPORT LIBRARIES AND MODULES #7#------------------------------------------------------------------------------#8from itertools import islice9import cv2, os, spaces, time, uuid10import gradio as gr11import numpy as np12 13from collections import defaultdict14from pathlib import Path15from PIL import Image16from typing import Iterator, Tuple, Union17from ultralytics import YOLO18from assets.ocean import Ocean19 20#--- Load the the weights ---#21WEIGHTS_PATH = 'models/weights/best.pt'22yolo = YOLO(WEIGHTS_PATH)23 24@spaces.GPU25def predict(source, conf_threshold, iou_threshold) -> Tuple[int, np.ndarray]:26 ''' This function predicts the results of the model on the given image.27 '''28 #--- Check if the source is a numpy array ---#29 if isinstance(source, np.ndarray): source = Image.fromarray(source)30 31 #--- Check if the source is a string ---#32 if isinstance(source, str): source = Image.open(source)33 34 #--- Check if the source is an image ---#35 if not isinstance(source, Image.Image): raise ValueError('The source must be an image.')36 37 #--- Predict the results ---#38 results = yolo.predict(source=source, conf=conf_threshold, iou=iou_threshold)39 40 #--- Draw the results on the image ---#41 for r in results:42 im_array = r.plot(line_width=3, labels=False)43 im = Image.fromarray(im_array[..., ::-1])44 45 return im, len(results[0].boxes.xywh)46 47@spaces.GPU48def track(video, conf_threshold, iou_threshold, analysis_fps = None) -> dict: # type: ignore49 ''' This function tracks the objects in the given video.50 51 Parameters52 ----------53 video : str54 The path to the video file.55 conf_threshold : float56 The confidence threshold for the detection.57 iou_threshold : float58 The IoU threshold for the detection.59 60 Returns61 -------62 dict63 The annotated frame and the count of the objects.64 '''65 def timestamps_frames_iterator(capture: cv2.VideoCapture) -> Iterator[np.ndarray]:66 ''' Iterate over the frames of a video capture and yield the timestamp and the frame.67 68 Parameters69 ----------70 capture : cv.VideoCapture71 The video capture object.72 73 Yields74 ------75 tuple76 The timestamp and the frame.77 '''78 #--- Iterate over the frames ---#79 while True:80 81 ret, frame = capture.read()82 83 #--- Return if no frame is read ---#84 if not ret: return85 86 #--- Yield the timestamp and the frame ---#87 yield np.asarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))88 89 global yolo90 91 #--- Load the video ---#92 cap = cv2.VideoCapture(video)93 assert cap.isOpened(), "Error: Cannot open the video file!!!" 94 print(f"FPS: {cap.get(cv2.CAP_PROP_FPS)}")95 96 #--- Set properties for the output video ---#97 video_codec = cv2.VideoWriter_fourcc(*'XVID') # Using avc1 codec for better compatibility98 fps = int(cap.get(cv2.CAP_PROP_FPS))99 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))100 height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))101 102 output_video_name = f".gradio/videos/{uuid.uuid4()}.avi"103 os.makedirs(os.path.dirname(output_video_name), exist_ok=True)104 output_video_writer = cv2.VideoWriter(output_video_name, video_codec, fps, (width, height))105 106 #--- Track history for each object ---#107 track_history = defaultdict(lambda: [])108 109 #--- Process each frame in the video ---#110 iterator = timestamps_frames_iterator(cap)111 if analysis_fps is not None:112 step = max(1, round(cap.get(cv2.CAP_PROP_FPS) / analysis_fps))113 iterator = islice(iterator, None, None, step)114 115 for frame in iterator:116 #--- Run YOLO detection and tracking ---#117 results = yolo.track(source=frame, persist=True, conf=conf_threshold, iou=iou_threshold)118 119 if results[0] is None or results[0].boxes.id is None:120 annotated_frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)121 count_number = 0122 else:123 #--- Extract detection results ---#124 boxes = results[0].boxes.xywh.cpu()125 track_ids = results[0].boxes.id.int().cpu().tolist()126 annotated_frame = results[0].plot()127 annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)128 129 #--- Draw tracking lines for each detected object ---#130 if len(boxes) > 0:131 for box, track_id in zip(boxes, track_ids):132 x, y, *_ = box133 track = track_history[track_id]134 # Store center point of bounding box135 track.append((float(x), float(y)))136 # Keep only the last 30 points for the trail137 if len(track) > 30: track.pop(0)138 139 # Draw the tracking trail140 points = np.array(track).reshape(-1, 1, 2).astype(np.int32)141 cv2.polylines(annotated_frame, [points], isClosed=False, color=(230, 230, 230), thickness=2)142 143 count_number = len(boxes)144 #--- Write the frame to the output video ---#145 output_video_writer.write(annotated_frame)146 147 #--- Yield the annotated frame ---#148 yield {'frame': annotated_frame[:, :, ::-1], 'count' : count_number}149 150 #--- Release the video capture and writer ---#151 cap.release()152 output_video_writer.release()153 track_history.clear()154 155 yolo = YOLO(WEIGHTS_PATH)156 print(f"Ouput Video File: ", output_video_name)157 158 #--- Return the output video ---#159 return {'frame': annotated_frame[:, :, ::-1], 'count' : count_number, 'video': output_video_name} 160 161#--- Create the interface ---#162theme = Ocean()163 164with gr.Blocks(165 css_paths=[Path("assets/ocean.css")], # Load custom CSS file166 theme=theme # Set the theme 167) as app:168 # #--- Welcome screen ---#169 # with gr.Group(visible=True) as welcome_screen:170 # #--- Include icons ---#171 # with gr.Column(elem_classes='welcome-container'):172 # gr.HTML("""173 # <video autoplay muted loop class="video-background">174 # <source src="https://videos.pexels.com/video-files/1918465/1918465-uhd_2560_1440_24fps.mp4" type="video/mp4">175 # Your browser does not support HTML5 video.176 # </video>177 # """)178 # with gr.Column(elem_classes="content-overlay"):179 # gr.HTML('<h1 class="welcome-title">YOLO-FishScale</h1>'\180 # '<p class="welcome-text">A real-time object detection and tracking system for marine life <br><br>By @Andrea Vincenzo Ricciardi</p>'\181 # '<div class="social-links">\182 # <a href="https://github.com/Andyvince01" target="_blank"><i class="fa fa-github"></i></a> \183 # <a href="https://www.linkedin.com/in/andrea-vincenzo-ricciardi-b50332262/ target="_blank"><i class="fa fa-linkedin"></i></a>'\184 # )185 # start_btn = gr.Button("Start", elem_classes="start-button", visible=True)186 187 #--- Main tabs ---#188 with gr.Tabs(visible=True) as main_tabs:189 #--- Image Detection Tab ---#190 with gr.TabItem("📸 Image Detection"):191 gr.HTML('<h1 class="video-title">Real-Time Fish Detection</h1>\192 <img src="https://www.diem.unisa.it/rescue/img/logo_standard.png" class="logo" style="width: 2.4cm; height: 2.4cm; overflow=hidden"/>\193 ')194 gr.HTML('<p class="video-instructions">Upload an image to detect fish within it.</p>')195 gr.Interface(196 fn=predict,197 inputs=[198 gr.Image(type="pil", label="Upload Image", value='https://cdn.pixabay.com/photo/2017/05/26/23/35/underwater-2347255_1280.jpg'), # Input image199 gr.Slider(minimum=0, maximum=1, value=0.15, label="Confidence threshold"), # Confidence threshold200 gr.Slider(minimum=0, maximum=1, value=0.6, label="IoU threshold") # IoU threshold201 ],202 outputs=[203 gr.Image(type='pil', label="Processed Image"), # Processed image204 gr.Number(label="Fish Count", value=0, visible=True) # Fish count205 ]206 )207 208 #--- Video Tracking Tab ---#209 with gr.TabItem("🎥 Video Tracking"):210 gr.HTML('<h1 class="video-title">Real-Time Video Tracking</h1>\211 <img src="https://www.diem.unisa.it/rescue/img/logo_standard.png" class="logo" style="width: 2.4cm; height: 2.4cm; overflow=hidden"/>\212 ')213 with gr.Row():214 gr.HTML('<p class="video-instructions">Upload a video to track marine life in real-time.</p>')215 216 with gr.Row():217 with gr.Column():218 video_input = gr.Video(label="Upload Video", value='https://videos.pexels.com/video-files/2556839/2556839-hd_1920_1080_25fps.mp4')219 conf_slider = gr.Slider(minimum=0, maximum=1, value=0.5, label="Confidence threshold")220 iou_slider = gr.Slider(minimum=0, maximum=1, value=0.6, label="IoU threshold")221 with gr.Row():222 clear_btn = gr.ClearButton(components=[video_input, conf_slider, iou_slider])223 submit_btn = gr.Button("Submit", elem_classes="button", visible=True)224 225 with gr.Column():226 output_image = gr.Image(type='pil', label="Processed Video", streaming=True)227 count_number = gr.Number(label="Fish Count", value=0)228 download_btn = gr.DownloadButton("Download Video", visible=False)229 230 clear_btn.add([output_image, count_number])231 232 #--- Set the track wrapper ---#233 def track_wrapper(video, conf_threshold, iou_threshold):234 #--- Initialize the generator ---#235 generator = track(video, conf_threshold, iou_threshold)236 237 #--- Iterate over the generator ---#238 while True:239 try:240 results = next(generator)241 yield (results['frame'], results['count'], gr.update(visible=False))242 #--- Stop the iteration if the video is processed ---#243 except StopIteration as final:244 # Check if the video was processed245 if final.value.get('video', None) is not None:246 yield (final.value['frame'], final.value['count'], gr.update(visible=True, value=final.value['video'])) 247 break248 249 #--- Set the button callbacks ---# 250 # start_btn.click(251 # fn=lambda: (gr.Group(visible=False), gr.Tabs(visible=True)),252 # outputs=[welcome_screen, main_tabs]253 # )254 255 submit_btn.click(256 fn=track_wrapper,257 inputs=[video_input, conf_slider, iou_slider],258 outputs=[output_image, count_number, download_btn]259 )260 261if __name__ == "__main__":262 app.launch()