add crack split.py

.idea/workspace.xml

@@ -4,10 +4,12 @@
     <option name="autoReloadType" value="SELECTIVE" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="d7806539-b6d6-42e7-bb45-1565f5d54891" name="Changes" comment="fix runtime error in coco evaluator">
+    <list default="true" id="d7806539-b6d6-42e7-bb45-1565f5d54891" name="Changes" comment="add crack split.py">
+      <change afterPath="$PROJECT_DIR$/models/tools/draw.py" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/demo.py" beforeDir="false" afterPath="$PROJECT_DIR$/demo.py" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/yolo_dataset.py" beforeDir="false" afterPath="$PROJECT_DIR$/yolo_dataset.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/models/yolo_crack.py" beforeDir="false" afterPath="$PROJECT_DIR$/models/yolo_crack.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/yolo_model.py" beforeDir="false" afterPath="$PROJECT_DIR$/yolo_model.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -56,37 +58,41 @@
     <option name="hideEmptyMiddlePackages" value="true" />
     <option name="showLibraryContents" value="true" />
   </component>
-  <component name="PropertiesComponent">{
-  &quot;keyToString&quot;: {
-    &quot;Python.demo.executor&quot;: &quot;Debug&quot;,
-    &quot;Python.evaluator.executor&quot;: &quot;Debug&quot;,
-    &quot;Python.extract.executor&quot;: &quot;Run&quot;,
-    &quot;Python.streamlit.executor&quot;: &quot;Debug&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": "Debug",
+    "Python.evaluator.executor": "Debug",
+    "Python.extract.executor": "Run",
+    "Python.streamlit.executor": "Run",
+    "Python.yolo_crack.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/tools",
+    "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$/tools" />
+      <recent name="$PROJECT_DIR$/datasets" />
       <recent name="$PROJECT_DIR$/tests" />
       <recent name="$PROJECT_DIR$" />
     </key>
     <key name="MoveFile.RECENT_KEYS">
+      <recent name="$PROJECT_DIR$/models" />
       <recent name="$PROJECT_DIR$/tests" />
       <recent name="$PROJECT_DIR$" />
     </key>
   </component>
-  <component name="RunManager" selected="Python.streamlit">
+  <component name="RunManager" selected="Python.yolo_crack">
     <configuration name="demo" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
       <module name="detector" />
       <option name="ENV_FILES" value="" />
@@ -157,6 +163,29 @@
       <option name="INPUT_FILE" value="" />
       <method v="2" />
     </configuration>
+    <configuration name="yolo_crack" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="detector" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/models/yolo_crack.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
     <configuration name="yolo_dataset" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
       <module name="detector" />
       <option name="ENV_FILES" value="" />
@@ -199,12 +228,14 @@
     <list>
       <item itemvalue="FastAPI.detector" />
       <item itemvalue="Python.streamlit" />
+      <item itemvalue="Python.yolo_crack" />
       <item itemvalue="Python.demo" />
       <item itemvalue="Python.evaluator" />
       <item itemvalue="Python.yolo_dataset" />
     </list>
     <recent_temporary>
       <list>
+        <item itemvalue="Python.yolo_crack" />
         <item itemvalue="Python.demo" />
         <item itemvalue="Python.evaluator" />
         <item itemvalue="Python.yolo_dataset" />
@@ -239,7 +270,8 @@
       <workItem from="1706274709532" duration="27000" />
       <workItem from="1706445222599" duration="1814000" />
       <workItem from="1706449264770" duration="4974000" />
-      <workItem from="1706517295469" duration="1726000" />
+      <workItem from="1706517295469" duration="2595000" />
+      <workItem from="1706855634459" duration="14109000" />
     </task>
     <task id="LOCAL-00001" summary="init commit">
       <option name="closed" value="true" />
@@ -473,7 +505,23 @@
       <option name="project" value="LOCAL" />
       <updated>1706517885858</updated>
     </task>
-    <option name="localTasksCounter" value="30" />
+    <task id="LOCAL-00030" summary="filter out macos cache files">
+      <option name="closed" value="true" />
+      <created>1706519049182</created>
+      <option name="number" value="00030" />
+      <option name="presentableId" value="LOCAL-00030" />
+      <option name="project" value="LOCAL" />
+      <updated>1706519049183</updated>
+    </task>
+    <task id="LOCAL-00031" summary="add crack split.py">
+      <option name="closed" value="true" />
+      <created>1706866629031</created>
+      <option name="number" value="00031" />
+      <option name="presentableId" value="LOCAL-00031" />
+      <option name="project" value="LOCAL" />
+      <updated>1706866629031</updated>
+    </task>
+    <option name="localTasksCounter" value="32" />
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
@@ -518,7 +566,9 @@
     <MESSAGE value="IoU evaluator" />
     <MESSAGE value="Update UI" />
     <MESSAGE value="fix runtime error in coco evaluator" />
-    <option name="LAST_COMMIT_MESSAGE" value="fix runtime error in coco evaluator" />
+    <MESSAGE value="filter out macos cache files" />
+    <MESSAGE value="add crack split.py" />
+    <option name="LAST_COMMIT_MESSAGE" value="add crack split.py" />
   </component>
   <component name="XDebuggerManager">
     <breakpoint-manager>
@@ -545,7 +595,8 @@
     <SUITE FILE_PATH="coverage/detector$yolo_dataset.coverage" NAME="yolo_dataset Coverage Results" MODIFIED="1705852113469" 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="1706107083258" 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="1706108414052" 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$streamlit.coverage" NAME="streamlit Coverage Results" MODIFIED="1706518941359" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="" />
+    <SUITE FILE_PATH="coverage/detector$streamlit.coverage" NAME="streamlit Coverage Results" MODIFIED="1706949799318" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="" />
+    <SUITE FILE_PATH="coverage/detector$yolo_crack.coverage" NAME="yolo_crack Coverage Results" MODIFIED="1706950034213" 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>

demo.py

@@ -4,6 +4,7 @@ import time
 from functools import wraps
 from io import StringIO
 from zipfile import ZipFile
+from tempfile import mktemp
 
 import streamlit as st
 from PIL import Image
@@ -11,9 +12,10 @@ from PIL import Image
 import evaluator
 from yolo_dataset import YoloDataset
 from yolo_model import YoloModel
+from models.yolo_crack import YoloModel as CrackModel
 
 fire_and_smoke = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
-crack = YoloModel("SHOU-ISD/yolo-cracks", "best.pt")
+crack = CrackModel("SHOU-ISD/yolo-cracks", "last4.pt", "SHOU-ISD/yolo-cracks", "best.pt")
 coco = YoloModel("ultralyticsplus/yolov8s", "yolov8s.pt")
 
 
@@ -71,15 +73,15 @@ def detect(model: YoloModel):
 
     if buffer:
         # Object Detecting
-        with st.spinner('Wait for it...'):
+        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)
+            filename = mktemp(suffix=buffer.name)
+            Image.open(buffer).save(filename)
+            res_img = model.preview_detect(filename, confidence / 100.0)
             t2 = time.time()
 
         # Displaying the image

models/__init__.py

@@ -0,0 +1 @@
+from .yolo_crack import *

models/tools/__init__.py

@@ -0,0 +1 @@
+from .split import *

models/tools/draw.py

@@ -0,0 +1,65 @@
+import random
+from PIL import ImageDraw
+
+
+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))
+    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
+    img_draw = ImageDraw.Draw(img)
+    img_draw.rectangle((c1[0], c1[1], c2[0], c2[1]), outline=color, width=tl)
+    if label:
+        tf = max(tl - 1, 1)  # font thickness
+        x1, y1, x2, y2 = img_draw.textbbox(c1, label, stroke_width=tf)
+        img_draw.rectangle((x1, y1, x2, y2), fill=color)
+        img_draw.text((x1, y1), label, fill=(255, 255, 255))
+
+
+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()
+        if conf < confidence:
+            continue
+        [rect] = box.xyxy.tolist()
+        text = f'{result.names[cl]}: {conf: 0.2f}'
+        plot_one_box(x=rect, img=pil_img, label=text)
+
+    return pil_img
+
+
+def add_bboxes2(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']
+        conf = box['conf']
+        if conf < confidence:
+            continue
+        rect = box['xyxy']
+        text = f'{cl}: {conf: 0.2f}'
+        plot_one_box(x=rect, img=pil_img, label=text)
+
+    return pil_img

models/tools/split.py

@@ -0,0 +1,175 @@
+import json
+import copy
+from typing import Optional
+from PIL import Image
+
+bbox = [float, float, float, float]
+
+annotation = {
+    "id": int,
+    "image_id": int,
+    "category_id": int,
+    "bbox": bbox,
+    "ignore": int,
+    "iscrowd": int,
+    "area": float,
+}
+
+small_image = {
+    "image": Image,
+    "area": bbox
+}
+
+
+def split_image(image: Image,
+                hint_size_min: tuple[int, int],
+                hint_size_max: tuple[int, int],
+                overlap: float = 0.1) -> list[small_image]:
+    """
+    Given an image and a hint size, split the image into a list of images.
+    New images are overlapped with other images by the overlap ratio.
+    :param image: The image to split. typically a large image. 1kx1k ~ 10kx10k
+    :param hint_size_min: The minimum size of the output image.
+    :param hint_size_max: The maximum size of the output image.
+    :param overlap: The overlap ratio of the output image.
+    :return: A list of images.
+    """
+    Wi, Hi = image.size
+    Wmin, Hmin = hint_size_min
+    Wmax, Hmax = hint_size_max
+    assert Wmin <= Wmax <= Wi
+    assert Hmin <= Hmax <= Hi
+    w_search = search(Wi, Wmin, Wmax, overlap)
+    h_search = search(Hi, Hmin, Hmax, overlap)
+    if w_search is None or h_search is None:
+        raise ValueError('The image is too small to split.')
+    w_count, output_width, last_output_width, width_overlap = w_search
+    h_count, output_height, last_output_height, height_overlap = h_search
+    images = []
+    for h_index in range(h_count):
+        h = h_index * (output_height - height_overlap)
+        for w_index in range(w_count):
+            w = w_index * (output_width - width_overlap)
+            small = {
+                "image": image.crop((w, h, w + output_width, h + output_height)),
+                "area": (w, h, output_width, output_height)
+            }
+            images.append(small)
+        if last_output_width > 0:
+            w = Wi - output_width
+            small = {
+                "image": image.crop((w, h, w + output_width, h + output_height)),
+                "area": (w, h, output_width, output_height)
+            }
+            images.append(small)
+    return images
+
+
+def search(input: int,
+           output_min: int,
+           output_max: int,
+           overlap: float) -> Optional[tuple[int, int, int, int]]:
+    """
+    example 1:
+    input: 8000, output: 1000, overlap: 0.1
+    8000 // (1000 - 100) = 8
+    8000 % (1000 - 100) = 800
+    count = 8, output = 1000, last_output = 800, overlap_pixels = 100
+
+    example 2:
+    input: 7200, output: 800, overlap: 0.1
+    7200 // (800 - 80) = 10
+    7200 % (800 - 80) = 0
+    count = 10, output = 800, last_output = 0, overlap_pixels = 80
+
+    :param input: The length of the input image.
+    :param output_min: The minimum length of the output image.
+    :param output_max: The maximum length of the output image.
+    :param overlap: The overlap ratio of the output image.
+    :return: A tuple of (count, output, last_output, overlap_pixels).
+    """
+
+    for output in range(output_max, output_min - 1, -1):
+        overlap_pixels = int(output * overlap)
+        last_output = input % (output - overlap_pixels)
+        if last_output == 0 or output_min <= last_output <= output_max:
+            count = input // (output - overlap_pixels)
+            return count, output, last_output, overlap_pixels
+
+    return None
+
+
+def box_intersected(box1: bbox, box2: bbox) -> bool:
+    """
+    Check if two boxes are intersected.
+    :param box1: The first box.
+    :param box2: The second box.
+    :return: True if the two boxes are intersected.
+    """
+    x1, y1, w1, h1 = box1
+    x2, y2, w2, h2 = box2
+    return x1 < x2 + w2 and x2 < x1 + w1 and y1 < y2 + h2 and y2 < y1 + h1
+
+
+def fit_in_area(annotations: list[annotation], in_area: bbox) -> list[annotation]:
+    result = []
+    for old in annotations:
+        ann = copy.deepcopy(old)
+        result.append(ann)
+        x, y, w, h = ann["bbox"]
+        if x < in_area[0]:
+            ann["bbox"][0] = 0
+        else:
+            ann["bbox"][0] -= in_area[0]
+        if y < in_area[1]:
+            ann["bbox"][1] = 0
+        else:
+            ann["bbox"][1] -= in_area[1]
+        if x + w > in_area[0] + in_area[2]:
+            ann["bbox"][2] = in_area[2] - ann["bbox"][0]
+        if y + h > in_area[1] + in_area[3]:
+            ann["bbox"][3] = in_area[3] - ann["bbox"][1]
+    return result
+
+
+small_image_with_labels = {
+    "image": Image,
+    "area": bbox,
+    "labels": list[annotation]
+}
+
+
+def split_image_with_labels(image: Image,
+                            labels: list[annotation],
+                            hint_size_min: tuple[int, int],
+                            hint_size_max: tuple[int, int],
+                            overlap: float = 0.1) -> list[small_image]:
+    small_imgs = split_image(image, hint_size_min, hint_size_max, overlap)
+    result = []
+    for small_img in small_imgs:
+        small_labels = [ann for ann in labels if box_intersected(ann["bbox"], small_img["area"])]
+        small_labels = fit_in_area(small_labels, small_img["area"])
+        result.append({
+            "image": small_img["image"],
+            "area": small_img["area"],
+            "labels": small_labels
+        })
+    return result
+
+
+def main():
+    image = Image.open('../datasets/Das3300161.jpg')
+    small_imgs = split_image(image, (800, 800), (1000, 1000), 0.1)
+    labels = json.load(open('../datasets/result.json'))
+    annotations = list(filter(lambda ann: ann["image_id"] == 28, labels["annotations"]))
+    for small_img in small_imgs:
+        small_labels = [ann for ann in annotations if box_intersected(ann["bbox"], small_img["area"])]
+        small_labels = fit_in_area(small_labels, small_img["area"])
+        # save small_labels to json
+        json.dump(small_labels, open('datasets/' + str(small_img["area"]) + '.json', 'w'))
+        # save small_image["image"] to file
+        small_img["image"].save('datasets/' + str(small_img["area"]) + '.jpg')
+
+
+if __name__ == '__main__':
+    main()

models/yolo_crack.py

@@ -0,0 +1,295 @@
+from models.tools import split
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+
+import torch
+import cv2
+import numpy as np
+import math
+
+from models.tools.draw import add_bboxes2
+
+
+class YoloModel:
+    def __init__(self, seg_repo_name: str, seg_file_name: str, det_repo_name: str, det_file_name: str):
+        seg_weight_file = YoloModel.download_weight_file(seg_repo_name, seg_file_name)
+        det_weight_file = YoloModel.download_weight_file(det_repo_name, det_file_name)
+        self.seg_model = YOLO(seg_weight_file)
+        self.det_model = YOLO(det_weight_file)
+
+    @staticmethod
+    def download_weight_file(repo_name: str, file_name: str):
+        return hf_hub_download(repo_name, file_name)
+
+    def preview_detect(self, im, confidence):
+        results = self.detect(im)
+        res_img = Image.open(im)
+        res = {
+            'boxes': [
+                {
+                    'xyxy': [x1, y1, x2, y2],
+                    'cls': cls,
+                    'conf': conf
+                } for x1, y1, x2, y2, conf, cls in results
+            ]
+        }
+        res_img = add_bboxes2(res_img, res, confidence)
+        return res_img
+
+    def detect(self, source):
+        pred_bbox_list = []  # 初始化该图像bbox列表
+        threshold = 50  # 暂定bbox merge 阈值为50, 后期可根据用户需求做自适应调整
+        strategy = "distance"  # 暂定bbox merge 策略为distance
+
+        seg_img_list = self._seg_ori_img(source)  # 对该图像进行路面分割
+        assert len(seg_img_list) == 1, "seg_img_list out of range"
+        road_img = Image.fromarray(cv2.cvtColor(seg_img_list[0], cv2.COLOR_BGR2RGB))
+        small_imgs = split.split_image(road_img, (640, 640), (1080, 1080), 0.1)  # 对路面图像进行小图分割
+        num = 0
+        for small_img in small_imgs:
+            num += 1
+            results = self.det_model(source=small_img["image"])
+            for result in results:
+                temp_bbox_list = result.boxes.xyxy  # 获取检测结果中的bbox坐标（此处使用xyxy格式）
+                w_bias = small_img["area"][0]
+                h_bias = small_img["area"][1]
+                temp_bbox_list = self._bbox_map(temp_bbox_list, w_bias, h_bias)  # 将bbox坐标映射到原始大图坐标系中
+                temp_bbox_cls = result.boxes.cls  # 获取检测结果中的class
+                temp_bbox_conf = result.boxes.conf  # 获取检测结果中的confidence
+                assert len(temp_bbox_list) == len(temp_bbox_cls) == len(
+                    temp_bbox_conf), 'different number of matrix size'
+                for i in range(len(temp_bbox_list)):  # 整合bbox、conf和class到一个数组中
+                    temp_bbox_list[i].append(temp_bbox_conf[i])
+                    temp_bbox_list[i].append(temp_bbox_cls[i])
+                pred_bbox_list += temp_bbox_list  # 将单张大图分割后的全体小图得到的检测结果（bbox、conf、class）整合到一个list
+        pred_bbox_list = self._merge_box(pred_bbox_list, threshold, strategy=strategy)  # 调用指定算法，对bbox进行分析合并
+
+        return pred_bbox_list
+
+    def _seg_ori_img(self, source):
+        """
+        分割原始图像中的沥青路面区域
+        :param source: 图像路径
+        :return: 分割得到的沥青路面图像（尺寸与原始图像一致，非路面区域用白色填充）
+        """
+        ori_img = cv2.imread(source)
+        ori_size = ori_img.shape
+        results = self.seg_model(source=source)
+        seg_img_list = []
+
+        for result in results:
+            if result.masks is not None and len(result.masks) > 0:  # 检测到路面时
+                masks_data = result.masks.data
+                obj_masks = masks_data[:]
+                road_mask = torch.any(obj_masks, dim=0).int() * 255
+                mask = road_mask.cpu().numpy()
+                Mask = mask.astype(np.uint8)
+                mask_res = cv2.resize(Mask, (ori_size[1], ori_size[0]), interpolation=cv2.INTER_CUBIC)
+
+            else:  # 检测不到路面时保存纯黑色图像
+                mask_res = np.zeros((ori_size[0], ori_size[1], 3), dtype=np.uint8)
+
+            mask_region = mask_res == 0
+            ori_img[mask_region] = 255  # 判断条件置0掩码为黑，置255背景为白
+            seg_img_list.append(ori_img)
+
+        return seg_img_list
+
+    def _bbox_map(self, bbox_list, w, h):
+        """
+        将小图中的bbox坐标映射到原始图像中
+        :param bbox_list: 小图中的bbox数组
+        :param w: 小图在原始图像中的偏置w
+        :param h: 小图在原始图像中的偏置h
+        :return: 该bbox数组在原始图像中的坐标
+        """
+        if isinstance(bbox_list, torch.Tensor):
+            bbox_list = bbox_list.tolist()
+        for bbox in bbox_list:
+            bbox[0] += w
+            bbox[1] += h
+            bbox[2] += w
+            bbox[3] += h
+
+        return bbox_list
+
+    def _xywh2xyxy(self, box_list):
+        """
+        YOLO标签，xywh转xyxy
+        :param box_list: bbox数组（xywh）
+        :return: bbox数组（xyxy）
+        """
+        new_box_list = []
+        for box in box_list:
+            x1 = box[0] - box[2] / 2
+            y1 = box[1] - box[3] / 2
+            x2 = box[0] + box[2] / 2
+            y2 = box[1] + box[3] / 2
+            new_box_list.append([x1, y1, x2, y2])
+
+        return new_box_list
+
+    def _xyxy2xywh(self, box_list):
+        """
+        YOLO标签，xyxy转xywh
+        :param box_list: bbox数组（xyxy）
+        :return: bbox数组（xywh）
+        """
+        new_box_list = []
+        for box in box_list:
+            x1 = (box[0] + box[2]) / 2
+            y1 = (box[1] + box[3]) / 2
+            w = (box[2] - box[0])
+            h = (box[3] - box[1])
+            new_box_list.append([x1, y1, w, h])
+
+        return new_box_list
+
+    def _nor2std(self, box_list, img_w, img_h):
+        """
+        YOLO标签，标准化坐标映射到原始图像
+        :param box_list: bbox数组（nor）
+        :param img_w: 原始图像宽度
+        :param img_h: 原始图像高度
+        :return: bbox数组（在原始图像中的坐标）
+        """
+        for box in box_list:
+            box[0] *= img_w
+            box[1] *= img_h
+            box[2] *= img_w
+            box[3] *= img_h
+
+    def _std2nor(self, box_list, img_w, img_h):
+        """
+        YOLO标签，原始图像坐标转标准化坐标
+        :param box_list: bbox数组（std）
+        :param img_w: 原始图像宽度
+        :param img_h: 原始图像高度
+        :return: bbox数组（标准化坐标）
+        """
+        for box in box_list:
+            box[0] /= img_w
+            box[1] /= img_h
+            box[2] /= img_w
+            box[3] /= img_h
+
+    def _judge_merge_by_center_distance(self, center_box1, center_box2, distance_threshold):
+        """
+        根据bbox中心坐标间距，判断是否进行bbox合并
+        :param center_box1: box1的中心坐标
+        :param center_box2: box2的中心坐标
+        :param distance_threshold: 间距阈值
+        :return: 若间距小于阈值，进行合并（Ture）；反之则忽略（False）
+        """
+        x1 = center_box1[0]
+        x2 = center_box2[0]
+        y1 = center_box1[1]
+        y2 = center_box2[1]
+        distance = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+        if distance < distance_threshold:
+            return True
+        else:
+            return False
+
+    def _judge_merge_by_overlap_area(self, std_box1, std_box2, overlap_threshold):
+        """
+        根据bbox交叉面积，判断是否进行bbox合并
+        :param std_box1: box1的标准坐标
+        :param std_box2: box2的标准坐标
+        :param overlap_threshold: 交叉面积阈值
+        :return: 若交叉面积大于阈值，进行合并（True）；反之则忽略（False）
+        """
+        x1 = max(std_box1[0], std_box2[0])
+        y1 = max(std_box1[1], std_box2[1])
+        x2 = min(std_box1[2], std_box2[2])
+        y2 = min(std_box1[3], std_box2[3])
+        width = max(0, x2 - x1)
+        height = max(0, y2 - y1)
+        area = width * height
+        if area < overlap_threshold:
+            return False
+        else:
+            return True
+
+    def _basic_merge(self, box1, box2):
+        """
+        合并两个box，生成新的box坐标
+        :param box1: box1坐标（std）
+        :param box2: box2坐标（std）
+        :return: 新box坐标（std）
+        """
+        x11 = box1[0]
+        y11 = box1[1]
+        x12 = box1[2]
+        y12 = box1[3]
+        x21 = box2[0]
+        y21 = box2[1]
+        x22 = box2[2]
+        y22 = box2[3]
+        new_x1 = min(x11, x12, x21, x22)
+        new_y1 = min(y11, y12, y21, y22)
+        new_x2 = max(x11, x12, x21, x22)
+        new_y2 = max(y11, y12, y21, y22)
+        assert len(box1) == len(box2), 'box1 and box2 has different size'
+        if len(box1) == 6:  # 此时，box中带有conf和class，其结构为[x1, y1, x2, y2, conf, class]
+            avg_conf = (box1[4] + box2[4]) / 2
+            clas = box1[5]
+            new_box = [new_x1, new_y1, new_x2, new_y2, avg_conf, clas]
+        else:
+            new_box = [new_x1, new_y1, new_x2, new_y2]
+
+        return new_box
+
+    def _update_list(self, bbox_list, del_index):
+        """
+        更新bbox数组，删除特定的bbox元素（已经被合并到其他box中的bbox）
+        :param bbox_list: bbox数组
+        :param del_index: 待删除bbox元素的rank
+        :return: 更新后的bbox数组
+        """
+        assert len(bbox_list) > del_index >= 0, 'del_index out of boundary'
+        bbox_list[del_index] = bbox_list[-1:][0]
+        bbox_list.pop()
+        return bbox_list
+
+    def _merge_box(self, std_bbox_list, threshold, strategy='overlap'):
+        """
+        bbox合并算法，根据选定的合并策略及阈值，进行bbox合并
+        :param std_bbox_list: std_bbox_list可有两种格式：(Array[N, 4] -> [x1, y1, x2, y2]; Array[N, 6] -> [x1, y1, x2, y2, conf, class])
+        :param threshold: 阈值
+        :param strategy: 合并策略（distance/overlap）
+        """
+        if isinstance(std_bbox_list, torch.Tensor):
+            std_bbox_list = std_bbox_list.tolist()
+        center_bbox_list = self._xyxy2xywh(std_bbox_list)
+        i = 0
+        while i < len(std_bbox_list):
+            j = i + 1
+            while j < len(std_bbox_list):
+                if strategy == 'overlap':
+                    assert i < len(std_bbox_list) and j < len(std_bbox_list), f'len={len(std_bbox_list)}, j={j}, i={i}'
+                    if self._judge_merge_by_overlap_area(std_bbox_list[i], std_bbox_list[j], threshold):
+                        std_bbox_list[i] = self._basic_merge(std_bbox_list[i], std_bbox_list[j])
+                        self._update_list(std_bbox_list, j)
+                        self._update_list(center_bbox_list, j)
+                        continue
+                else:
+                    if self._judge_merge_by_center_distance(center_bbox_list[i], center_bbox_list[j], threshold):
+                        std_bbox_list[i] = self._basic_merge(std_bbox_list[i], std_bbox_list[j])
+                        self._update_list(std_bbox_list, j)
+                        self._update_list(center_bbox_list, j)
+                        continue
+                j += 1
+            i += 1
+
+        return std_bbox_list
+
+
+def main():
+    model = YoloModel("SHOU-ISD/yolo-cracks", "last4.pt", "SHOU-ISD/yolo-cracks", "best.pt")
+    model.preview_detect('./datasets/Das1100209.jpg', 0.4).show()
+
+
+if __name__ == '__main__':
+    main()

yolo_model.py

@@ -1,50 +1,8 @@
-import random
-from PIL import ImageDraw, Image
+from PIL import Image
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 
-
-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))
-    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
-    img_draw = ImageDraw.Draw(img)
-    img_draw.rectangle((c1[0], c1[1], c2[0], c2[1]), outline=color, width=tl)
-    if label:
-        tf = max(tl - 1, 1)  # font thickness
-        x1, y1, x2, y2 = img_draw.textbbox(c1, label, stroke_width=tf)
-        img_draw.rectangle((x1, y1, x2, y2), fill=color)
-        img_draw.text((x1, y1), label, fill=(255, 255, 255))
-
-
-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()
-        if conf < confidence:
-            continue
-        [rect] = box.xyxy.tolist()
-        text = f'{result.names[cl]}: {conf: 0.2f}'
-        plot_one_box(x=rect, img=pil_img, label=text)
-
-    return pil_img
+from models.tools.draw import add_bboxes
 
 
 class YoloModel:
@@ -59,9 +17,10 @@ class YoloModel:
     def detect(self, im):
         return self.model(source=im)
 
-    def preview_detect(self, im, confidence):
-        results = self.model(source=im)
-        res_img = im
+    def preview_detect(self, filename, confidence):
+        image = Image.open(filename)
+        results = self.model(source=image)
+        res_img = image
         for result in results:
             res_img = add_bboxes(res_img, result, confidence)
         return res_img