Update app.py

app.py

@@ -14,12 +14,13 @@ import cv2
 import os
 import PIL.Image as Image
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from ultralytics import ASSETS, YOLO
+import numpy as np
+from ultralytics import YOLO
 
 # load trained model
 model = YOLO("best.pt")
 
+# image inference function
 def predict_image(img, conf_threshold, iou_threshold):
     results = model.predict(
         source=img,
@@ -36,9 +37,12 @@ def predict_image(img, conf_threshold, iou_threshold):
 
     return im
 
-current_directory = "/home/user/app/image" 
+# directory for examples
+image_directory = "/home/user/app/image"
+video_directory = "/home/user/app/video"
 
-iface = gr.Interface(
+# interface gradio setting for image
+image_iface = gr.Interface(
     fn=predict_image,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
@@ -49,11 +53,58 @@ iface = gr.Interface(
     title="Fire Detection using YOLOv8n on Gradio",
     description="Upload images for inference. The Ultralytics YOLOv8n trained model is used for this.",
     examples=[
-        [os.path.join(current_directory, "fire_image_1.jpg"), 0.25, 0.45],
-        [os.path.join(current_directory, "fire_image_3.jpg"), 0.25, 0.45],
+        [os.path.join(image_directory, "fire_image_1.jpg"), 0.25, 0.45],
+        [os.path.join(image_directory, "fire_image_3.jpg"), 0.25, 0.45],
         
     ]
 )
 
+# convert PIL image objects to numpy arrays 
+def pil_to_cv2(pil_image):
+    open_cv_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+    return open_cv_image
+
+# process video, convert frame to PIL image 
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    processed_frames = []
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        # Model expects PIL Image format
+        pil_img = Image.fromarray(frame[..., ::-1])  # Convert BGR to RGB
+        result = model.predict(source=pil_img)
+        for r in result:
+            im_array = r.plot()
+            processed_frames.append(Image.fromarray(im_array[..., ::-1]))  # Convert RGB back to BGR
+    cap.release()
+    # You may choose to display each frame or compile them back using cv2 or a similar library
+    # Display the processed frames
+    for frame in processed_frames:
+        cv2.imshow(pil_to_cv2(frame))
+        if cv2.waitKey(25) & 0xFF == ord('q'):
+            break
+    cv2.destroyAllWindows()
+    # return processed_frames[-1]  # Example, returning the last processed frame
+
+# interface setting for video
+video_iface = gr.Interface(
+    fn=process_video,
+    inputs=[
+        gr.Video(label="Upload Video", interactive=True)
+    ],
+    outputs=gr.Video(label="Result"),
+    title="Fire Detection using YOLOv8n on Gradio",
+    description="Upload video for inference. The Ultralytics YOLOv8n trained model is used for inference.",
+    examples=[
+        [os.path.join(video_directory, "video_fire_1.jpg")],
+        [os.path.join(video_directory, "video_fire_2.jpg")],
+    ]
+)
+
+
+demo = gr.TabbedInterface([image_iface, video_iface], ["Image Inference", "Video Inference"])
+
 if __name__ == '__main__':
     iface.launch()