Project Files committed

Project_README.md

@@ -0,0 +1,9 @@
+# Fire Detection using YOLO Nas
+
+A streamlit app that uses a YOLO Nas fine-tuned model go predict given a image having a fire, where the fire is coming from. This could be very helpful in places were Fires are really not welcomed!
+
+- Dataset Used: https://universe.roboflow.com/-jwzpw/continuous_fire/
+- App link: https://huggingface.co/spaces/Harsh72AI/Fire-Detection
+
+
+Our Fire-detection model has acheived a  mAP:0.50 value of 0.82

Roboflow dataset.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72ceed891dd60e5df8e25028a73c0ad422dd59a53012400e5753d4610955e641
+size 15759390

app.py

@@ -0,0 +1,33 @@
+import streamlit as st
+from pipeline import predictPipeline
+
+
+st.title('Fire and Smoke detection')
+st.write('Detects Fire or/and Smoke in a Photo \nPowered by YOLO-Nas medium model')
+
+st.write('')
+
+detect_pipeline = predictPipeline()
+
+st.info('Fire and Smoke Detection model loaded successfully!')
+
+
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+
+if uploaded_file is not None:
+    
+    with st.container():
+        col1, col2 = st.columns([3, 3])
+        col1.header('Input Image')
+        col1.image(uploaded_file, caption='Uploaded Image', use_column_width=True)
+
+        col1.text('')
+        col1.text('')
+
+        if st.button('Detect'):
+            detections = detect_pipeline.detect(img_path=uploaded_file)
+            detections_img = detect_pipeline.drawDetections2Image(img_path=uploaded_file, detections=detections)
+
+            col2.header('Detections')
+            col2.image(detections_img, caption='Predictions by model', use_column_width=True)
+

pipeline.py

@@ -0,0 +1,38 @@
+from super_gradients.training import models
+import cv2 as cv
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as pl
+
+
+class predictPipeline():
+    def __init__(self) -> None:
+        
+        # Load model
+        self.model = models.get('yolo_nas_m',
+                                num_classes=1,
+                                checkpoint_path='yolo_nas_m_model.pth')
+    
+    def detect(self, img_path):
+        image = Image.open(img_path).convert('RGB')
+        img_array = np.array(image)
+
+        preds = self.model.predict(img_array, conf=0.5)[0].prediction
+        bboxes_coordinates = []
+        for idx, bbox in enumerate(preds.bboxes_xyxy):
+            bboxes_coordinates.append([int(num) for num in bbox] + [round(preds.confidence[idx]*100, 2)])
+        return bboxes_coordinates
+    
+    
+    def drawDetections2Image(self, img_path, detections):
+        img = Image.open(img_path).convert('RGB')
+        img = np.array(img)
+        for bbox in detections:
+            x1, y1, x2, y2, score = bbox
+            cv.rectangle(img, pt1=(x1, y1), pt2=(x2, y2), color=(0, 255, 0), thickness=2)
+            cv.putText(img, text=f'{score}%', org=(x1, y1-2), fontFace=cv.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0, 0, 255), lineType=cv.LINE_AA)
+        img_detections = np.array(img)
+        return img_detections
+
+
+

requirements.txt

@@ -0,0 +1,4 @@
+super-gradients
+torch
+numpy
+streamlit

yolo_nas_m_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce5c48348ef8e1f7e17b3aa0ff97d81b06b9811699103f8ab723cb6455abd0b8
+size 681045674