Update UI

.idea/workspace.xml

@@ -4,7 +4,15 @@
     <option name="autoReloadType" value="SELECTIVE" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="d7806539-b6d6-42e7-bb45-1565f5d54891" name="Changes" comment="IoU evaluator" />
+    <list default="true" id="d7806539-b6d6-42e7-bb45-1565f5d54891" name="Changes" comment="IoU evaluator">
+      <change afterPath="$PROJECT_DIR$/yolo_model.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/Dockerfile" beforeDir="false" afterPath="$PROJECT_DIR$/Dockerfile" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/demo.py" beforeDir="false" afterPath="$PROJECT_DIR$/demo.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/evaluator.py" beforeDir="false" afterPath="$PROJECT_DIR$/evaluator.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/yolo_dataset.py" beforeDir="false" afterPath="$PROJECT_DIR$/yolo_dataset.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/yolo_fire.py" beforeDir="false" />
+    </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
     <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
@@ -37,25 +45,25 @@
     <option name="hideEmptyMiddlePackages" value="true" />
     <option name="showLibraryContents" value="true" />
   </component>
-  <component name="PropertiesComponent">{
-  &quot;keyToString&quot;: {
-    &quot;Python.demo.executor&quot;: &quot;Run&quot;,
-    &quot;Python.evaluator.executor&quot;: &quot;Debug&quot;,
-    &quot;Python.extract.executor&quot;: &quot;Run&quot;,
-    &quot;Python.yolo_dataset.executor&quot;: &quot;Run&quot;,
-    &quot;RunOnceActivity.OpenProjectViewOnStart&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
-    &quot;git-widget-placeholder&quot;: &quot;main&quot;,
-    &quot;last_opened_file_path&quot;: &quot;/Users/famer.me/PycharmProjects/detector/tests&quot;,
-    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
-    &quot;node.js.detected.package.tslint&quot;: &quot;true&quot;,
-    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
-    &quot;node.js.selected.package.tslint&quot;: &quot;(autodetect)&quot;,
-    &quot;nodejs_package_manager_path&quot;: &quot;npm&quot;,
-    &quot;settings.editor.selected.configurable&quot;: &quot;settings.qodana&quot;,
-    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "Python.demo.executor": "Run",
+    "Python.evaluator.executor": "Run",
+    "Python.extract.executor": "Run",
+    "Python.yolo_dataset.executor": "Run",
+    "RunOnceActivity.OpenProjectViewOnStart": "true",
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "git-widget-placeholder": "main",
+    "last_opened_file_path": "/Users/famer.me/PycharmProjects/detector/tests",
+    "node.js.detected.package.eslint": "true",
+    "node.js.detected.package.tslint": "true",
+    "node.js.selected.package.eslint": "(autodetect)",
+    "node.js.selected.package.tslint": "(autodetect)",
+    "nodejs_package_manager_path": "npm",
+    "settings.editor.selected.configurable": "settings.qodana",
+    "vue.rearranger.settings.migration": "true"
   }
-}</component>
+}]]></component>
   <component name="RecentsManager">
     <key name="CopyFile.RECENT_KEYS">
       <recent name="$PROJECT_DIR$/tests" />
@@ -66,7 +74,7 @@
       <recent name="$PROJECT_DIR$" />
     </key>
   </component>
-  <component name="RunManager" selected="Python.evaluator">
+  <component name="RunManager" selected="Python.yolo_dataset">
     <configuration name="demo" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
       <module name="detector" />
       <option name="ENV_FILES" value="" />
@@ -153,8 +161,8 @@
     </configuration>
     <recent_temporary>
       <list>
-        <item itemvalue="Python.evaluator" />
         <item itemvalue="Python.yolo_dataset" />
+        <item itemvalue="Python.evaluator" />
         <item itemvalue="Python.demo" />
       </list>
     </recent_temporary>
@@ -177,7 +185,9 @@
       <workItem from="1702301554378" duration="4025000" />
       <workItem from="1702450051021" duration="10065000" />
       <workItem from="1702476484538" duration="9613000" />
-      <workItem from="1705757850315" duration="13343000" />
+      <workItem from="1705757850315" duration="14186000" />
+      <workItem from="1705812482464" duration="711000" />
+      <workItem from="1705845330999" duration="5355000" />
     </task>
     <task id="LOCAL-00001" summary="init commit">
       <option name="closed" value="true" />
@@ -347,7 +357,15 @@
       <option name="project" value="LOCAL" />
       <updated>1705770964768</updated>
     </task>
-    <option name="localTasksCounter" value="22" />
+    <task id="LOCAL-00022" summary="IoU evaluator">
+      <option name="closed" value="true" />
+      <created>1705771207294</created>
+      <option name="number" value="00022" />
+      <option name="presentableId" value="LOCAL-00022" />
+      <option name="project" value="LOCAL" />
+      <updated>1705771207294</updated>
+    </task>
+    <option name="localTasksCounter" value="23" />
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
@@ -393,8 +411,9 @@
     <option name="LAST_COMMIT_MESSAGE" value="IoU evaluator" />
   </component>
   <component name="com.intellij.coverage.CoverageDataManagerImpl">
+    <SUITE FILE_PATH="coverage/detector$yolo_dataset.coverage" NAME="yolo_dataset Coverage Results" MODIFIED="1705850662461" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
     <SUITE FILE_PATH="coverage/detector$demo.coverage" NAME="demo Coverage Results" MODIFIED="1702302380580" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
-    <SUITE FILE_PATH="coverage/detector$evaluator.coverage" NAME="evaluator Coverage Results" MODIFIED="1705768332357" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
+    <SUITE FILE_PATH="coverage/detector$evaluator.coverage" NAME="evaluator Coverage Results" MODIFIED="1705847678258" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
     <SUITE FILE_PATH="coverage/detector$extract.coverage" NAME="yolo_dataset Coverage Results" MODIFIED="1705764465837" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
   </component>
 </project>

Dockerfile

@@ -12,5 +12,7 @@ RUN chown -R user $HOME/.cache
 
 RUN pip install -r requirements.txt
 RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=true \
-    HF_TOKEN=$(cat /run/secrets/HF_TOKEN) python yolo_fire.py
+    HF_TOKEN=$(cat /run/secrets/HF_TOKEN) python yolo_model.py \
+    --weight_files=SHOU-ISD/fire-and-smoke:yolov8n.pt \
+    --weight_files=SHOU-ISD/yolo-cracks:best.pt
 ENTRYPOINT ["streamlit", "run", "demo.py"]

demo.py

@@ -9,52 +9,73 @@ from PIL import Image
 
 import evaluator
 from yolo_dataset import YoloDataset
-from yolo_fire import detect
+from yolo_model import YoloModel
+
+fire_and_smoke = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
+crack = YoloModel("SHOU-ISD/yolo-cracks", "best.pt")
 
 
 def main():
     # Header & Page Config.
     st.set_page_config(
-        page_title="Fire&Smoke Detection",
+        page_title=f"Detection",
         layout="centered")
-    st.title("Fire&Smoke Detection:")
+
+    model = None
+    model_choice = None
+    with st.sidebar:
+        model_choice = st.radio("Select Model", ["Fire&Smoke", "Crack"])
+        if model_choice == "Fire&Smoke":
+            model = fire_and_smoke
+        elif model_choice == "Crack":
+            model = crack
+
+    st.title(f"{model_choice} Detection:")
 
     detect_tab, evaluate_tab = st.tabs(["Detect", "Evaluate"])
 
     with evaluate_tab:
-        buffer = st.file_uploader("Upload your Yolo Dataset here", type=["zip"])
+        evaluate(model)
+    with detect_tab:
+        detect(model)
 
-        if buffer:
-            with st.spinner('Wait for it...'):
-                # Slider for changing confidence
-                confidence = st.slider('Confidence Threshold', 0, 100, 30)
-                yolo_dataset = YoloDataset.from_zip_file(ZipFile(buffer))
-                capture_output(evaluator.evaluate)(coco_gt=yolo_dataset.to_coco(), loader=yolo_dataset.load_image,
-                                                   confidence_threshold=confidence / 100.0)
 
-    with detect_tab:
-        # This will let you upload PNG, JPG & JPEG File
-        buffer = st.file_uploader("Upload your Image here", type=["jpg", "png", "jpeg"])
-
-        if buffer:
-            # Object Detecting
-            with st.spinner('Wait for it...'):
-                # Slider for changing confidence
-                confidence = st.slider('Confidence Threshold', 0, 100, 30)
-
-                # Calculating time for detection
-                t1 = time.time()
-                im = Image.open(buffer)
-                # im.save("saved_images/image.jpg")
-                res_img = detect(im, confidence / 100.0)
-                t2 = time.time()
-
-            # Displaying the image
-            st.image(res_img, use_column_width=True)
-
-            # Printing Time
-            st.write("\n")
-            st.write("Time taken: ", t2 - t1, "sec.")
+def evaluate(model: YoloModel):
+    buffer = st.file_uploader("Upload your Yolo Dataset here", type=["zip"])
+
+    if buffer:
+        with st.spinner('Wait for it...'):
+            # Slider for changing confidence
+            confidence = st.slider('Confidence Threshold', 0, 100, 30)
+            yolo_dataset = YoloDataset.from_zip_file(ZipFile(buffer))
+            capture_output(evaluator.evaluate)(model=model.model,
+                                               coco_gt=yolo_dataset.to_coco(), loader=yolo_dataset.load_image,
+                                               confidence_threshold=confidence / 100.0)
+
+
+def detect(model: YoloModel):
+    # This will let you upload PNG, JPG & JPEG File
+    buffer = st.file_uploader("Upload your Image here", type=["jpg", "png", "jpeg"])
+
+    if buffer:
+        # Object Detecting
+        with st.spinner('Wait for it...'):
+            # Slider for changing confidence
+            confidence = st.slider('Confidence Threshold', 0, 100, 30)
+
+            # Calculating time for detection
+            t1 = time.time()
+            im = Image.open(buffer)
+            # im.save("saved_images/image.jpg")
+            res_img = model.preview_detect(im, confidence / 100.0)
+            t2 = time.time()
+
+        # Displaying the image
+        st.image(res_img, use_column_width=True)
+
+        # Printing Time
+        st.write("\n")
+        st.write("Time taken: ", t2 - t1, "sec.")
 
 
 def capture_output(func):

evaluator.py

@@ -6,12 +6,12 @@ from pycocotools.coco import COCO
 from tqdm import tqdm
 
 from yolo_dataset import YoloDataset
-from yolo_fire import model
+from yolo_model import YoloModel
 
 image_loader = Callable[[str], Image]
 
 
-def evaluate(coco_gt: COCO, loader: image_loader, confidence_threshold=0.6):
+def evaluate(model: YoloModel, coco_gt: COCO, loader: image_loader, confidence_threshold=0.6):
     # initialize evaluator with ground truth (gt)
     evaluator = CocoEvaluator(coco_gt=coco_gt, iou_types=["bbox"])
 
@@ -19,9 +19,10 @@ def evaluate(coco_gt: COCO, loader: image_loader, confidence_threshold=0.6):
     for image_id, annotations in tqdm(coco_gt.imgToAnns.items()):
         # get the inputs
         image = coco_gt.imgs[image_id]
-        results = model(source=loader(image["file_name"]))
+        results = model.model(source=loader(image["file_name"]))
         for result in results:
-            coco_anns = yolo_boxes_to_coco_annotations(image_id, result.boxes, confidence_threshold=confidence_threshold)
+            coco_anns = yolo_boxes_to_coco_annotations(image_id, result.boxes,
+                                                       confidence_threshold=confidence_threshold)
             if len(coco_anns) == 0:
                 continue
             evaluator.update(coco_anns)
@@ -47,7 +48,19 @@ def yolo_boxes_to_coco_annotations(image_id: int, yolo_boxes, confidence_thresho
     ]
 
 
-if __name__ == '__main__':
-    yolo_dataset = YoloDataset.from_zip_file('tests/coco8.zip')
+def main():
+    yolo_dataset = YoloDataset.from_path('tests/coco8.zip')
     coco_gt = yolo_dataset.to_coco()
-    evaluate(coco_gt=coco_gt, loader=yolo_dataset.load_image, confidence_threshold=0.1)
+    model = YoloModel("ultralyticsplus/yolov8s", "yolov8s.pt")
+    # model = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
+    evaluate(model=model, coco_gt=coco_gt, loader=yolo_dataset.load_image, confidence_threshold=0.1)
+    # Validate the model
+    metrics = model.model.val()  # no arguments needed, dataset and settings remembered
+    print(metrics.box.map)  # map50-95
+    print(metrics.box.map50)  # map50
+    print(metrics.box.map75)  # map75
+    print(metrics.box.maps)  # a list contains map50-95 of each category
+
+
+if __name__ == '__main__':
+    main()

yolo_dataset.py

@@ -12,6 +12,9 @@
 #     ├── 1.txt
 #     ├── 2.txt
 #     └── ...
+import os
+import pathlib
+import shutil
 from datetime import datetime
 from typing import Optional
 from zipfile import ZipFile
@@ -198,8 +201,36 @@ class YoloDataset:
         coco_ds.createIndex()
         return coco_ds
 
+    def to_yaml(self):
+        """
+        :return: Yaml format in
+            <code>
+            path: ../datasets/coco128 # dataset root dir
+            val: images/train2017 # val images (relative to 'path') 128 images
+
+            # Classes
+            names:
+              0: person
+            </code>
+        """
+        try:
+            # recursively create dir
+            pathlib.Path(f'datasets/{self._zip_file.filename}').mkdir(parents=True, exist_ok=True)
+            self._zip_file.extractall(f'datasets/{self._zip_file.filename}')
+            with open(f'datasets/{self._zip_file.filename}/data.yaml', 'w') as f:
+                f.write(f'path: ../datasets/{self._zip_file.filename}\n')
+                f.write(f'val: images/train2017\n')
+                f.write(f'\n')
+                f.write(f'# Classes\n')
+                f.write(f'names:\n')
+                for i in range(len(self._classes)):
+                    f.write(f'  {i}: {self._classes[i]}\n')
+        finally:
+            shutil.rmtree(f'datasets/{self._zip_file.filename}')
+
 
 if __name__ == '__main__':
-    dataset = YoloDataset.from_zip_file('tests/coco8.zip')
+    dataset = YoloDataset.from_path('tests/coco8.zip')
     coco = dataset.to_coco()
     print(coco)
+    dataset.to_yaml()

yolo_fire.py → yolo_model.py

@@ -3,12 +3,17 @@ from PIL import ImageDraw, Image
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 
-# Load model
-weight_file = hf_hub_download("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
-model = YOLO(weight_file)  # pretrained YOLOv8n model
 
-# Helper Functions for Plotting BBoxes
 def plot_one_box(x, img, color=None, label=None, line_thickness=None):
+    """
+    Helper Functions for Plotting BBoxes
+    :param x:
+    :param img:
+    :param color:
+    :param label:
+    :param line_thickness:
+    :return:
+    """
     width, height = img.size
     tl = line_thickness or round(0.002 * (width + height) / 2) + 1  # line/font thickness
     color = color or (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
@@ -22,8 +27,14 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None):
         img_draw.text((x1, y1), label, fill=(255, 255, 255))
 
 
-# Ploting Bounding Box on img
 def add_bboxes(pil_img, result, confidence=0.6):
+    """
+    Plotting Bounding Box on img
+    :param pil_img:
+    :param result:
+    :param confidence:
+    :return:
+    """
     for box in result.boxes:
         [cl] = box.cls.tolist()
         [conf] = box.conf.tolist()
@@ -36,16 +47,61 @@ def add_bboxes(pil_img, result, confidence=0.6):
     return pil_img
 
 
-def detect(im, confidence):
-    results = model(source=im)
-    res_img = im
-    for result in results:
-        res_img = add_bboxes(res_img, result, confidence)
-    return res_img
+class YoloModel:
+    def __init__(self, repo_name: str, file_name: str):
+        weight_file = YoloModel.download_weight_file(repo_name, file_name)
+        self.model = YOLO(weight_file)
 
-if __name__ == '__main__':
+    @staticmethod
+    def download_weight_file(repo_name: str, file_name: str):
+        return hf_hub_download(repo_name, file_name)
+
+    def detect(self, im):
+        return self.model(source=im)
+
+    def preview_detect(self, im, confidence):
+        results = self.model(source=im)
+        res_img = im
+        for result in results:
+            res_img = add_bboxes(res_img, result, confidence)
+        return res_img
+
+
+def test():
+    model = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
     im = Image.open("./tests/fire1.jpg")
-    results = model(source=im)
+    results = model.model(source=im)
     for result in results:
-        im = add_bboxes(im, result)
-    im.show()
+        im = add_bboxes(im, result, confidence=0.1)
+        print(result.boxes)
+
+
+def argument_parser():
+    """
+    Argument Parser
+    :return: args
+    """
+    import argparse
+    parser = argparse.ArgumentParser(description='Help for YoloModel')
+    parser.add_argument('--test', '-t', action='store_true', help='Run test')
+    # list of repo_name&file_name
+    parser.add_argument('--weight_files', '-w', nargs='+', help='List of weight files')
+    return parser.parse_args()
+
+
+def pre_cache_weight_files(weight_files: list[str]):
+    """
+    Pre-cache weight files
+    :return: None
+    """
+    for weight_file in weight_files:
+        weight_file = weight_file.split(":")
+        YoloModel.download_weight_file(weight_file[0], weight_file[1])
+
+
+if __name__ == '__main__':
+    args = argument_parser()
+    if args.test:
+        test()
+    else:
+        pre_cache_weight_files(args.weight_files)