diff --git a/.gitignore b/.gitignore new file mode 100644 index 000000000..32b24f8d9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +env/ +*.pyc +__pycache__/ \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000..cb4d4dfae --- /dev/null +++ b/Dockerfile @@ -0,0 +1,31 @@ +# Выбор базового образа с предустановленным Python +FROM python:3.9-slim + +# Установка неинтерактивного режима для APT (автоматический выбор ответов по умолчанию) +ENV DEBIAN_FRONTEND=noninteractive + +# Настройка временной зоны (пример для часовой зоны Москва) +ENV TZ=Europe/Moscow +RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone + +# Обновление списка пакетов и установка необходимых системных зависимостей +RUN apt-get update && apt-get install -y \ + git \ + python3-opencv \ + # Удаление списков пакетов после установки для уменьшения размера образа + && rm -rf /var/lib/apt/lists/* + +# Копирование файла требований зависимостей Python в контейнер +COPY requirements.txt /workspace/requirements.txt + +# Установка зависимостей Python из файла требований +RUN pip install --no-cache-dir -r /workspace/requirements.txt + +# Копирование оставшегося исходного кода проекта в контейнер +COPY . /workspace + +# Установка рабочей директории +WORKDIR /workspace + +# Команда по умолчанию, запускающая оболочку bash +CMD ["bash"] diff --git a/PSGUI.py b/PSGUI.py new file mode 100644 index 000000000..b5d804bd1 --- /dev/null +++ b/PSGUI.py @@ -0,0 +1,187 @@ +import sys +from PyQt5 import QtWidgets, QtGui +from PyQt5.QtWidgets import ( + QFileDialog, + QLabel, + QVBoxLayout, + QHBoxLayout, + QPushButton, + QSlider, + QWidget, + QApplication, +) +from PyQt5.QtCore import Qt +from PIL import Image, ImageDraw, ImageFont +import zipfile +from pathlib import Path +import pandas as pd +from tqdm import tqdm +from detect_function_dual import detect_image_dual +from detect_function import detect_image + +# Переключатель для выбора функции детекции +use_dual_function = ( + True # Установите в False для использования стандартной функции детекции +) + +# Путь к весам модели +weight_path = ( + r"C:\Users\pasha\OneDrive\Рабочий стол\yolo_weights\yolo_word_detectino21.pt" +) + + +class MainApp(QWidget): + def __init__(self): + super().__init__() + + self.initUI() + + def initUI(self): + self.setWindowTitle("YOLO Detection") + self.setGeometry(100, 100, 800, 600) + + layout = QVBoxLayout() + + self.imageLabel = QLabel(self) + self.imageLabel.setAlignment(Qt.AlignCenter) + layout.addWidget(self.imageLabel) + + self.loadButton = QPushButton("Load Image or Zip", self) + self.loadButton.clicked.connect(self.loadFile) + layout.addWidget(self.loadButton) + + self.confSlider = QSlider(Qt.Horizontal) + self.confSlider.setRange(0, 100) + self.confSlider.setValue(25) + self.confSlider.setTickInterval(5) + self.confSlider.setTickPosition(QSlider.TicksBelow) + layout.addWidget(QLabel("Confidence Threshold")) + layout.addWidget(self.confSlider) + + self.iouSlider = QSlider(Qt.Horizontal) + self.iouSlider.setRange(0, 100) + self.iouSlider.setValue(45) + self.iouSlider.setTickInterval(5) + self.iouSlider.setTickPosition(QSlider.TicksBelow) + layout.addWidget(QLabel("IoU Threshold")) + layout.addWidget(self.iouSlider) + + self.processButton = QPushButton("Process", self) + self.processButton.clicked.connect(self.processFile) + layout.addWidget(self.processButton) + + self.setLayout(layout) + + def draw_boxes(self, image, results): + draw = ImageDraw.Draw(image) + font = ImageFont.load_default() + for result in results: + x1, y1, x2, y2 = result["bbox"] + confidence = result["conf"] + class_id = result["cls"] + draw.rectangle([x1, y1, x2, y2], outline="red", width=2) + text = f"Class {int(class_id)} {confidence:.2f}" + text_size = draw.textbbox((0, 0), text, font=font) + text_location = [x1, y1 - text_size[3]] + if text_location[1] < 0: + text_location[1] = y1 + text_size[3] + draw.rectangle([x1, y1 - text_size[3], x1 + text_size[2], y1], fill="red") + draw.text((x1, y1 - text_size[3]), text, fill="white", font=font) + return image + + def loadFile(self): + options = QFileDialog.Options() + fileName, _ = QFileDialog.getOpenFileName( + self, + "Open Image or Zip", + "", + "Images (*.png *.xpm *.jpg);;Zip Files (*.zip)", + options=options, + ) + if fileName: + self.filePath = fileName + self.imageLabel.setPixmap(QtGui.QPixmap(fileName)) + + def processFile(self): + if not hasattr(self, "filePath"): + QtWidgets.QMessageBox.warning(self, "Error", "No file loaded") + return + + conf_thres = self.confSlider.value() / 100 + iou_thres = self.iouSlider.value() / 100 + file_path = Path(self.filePath) + detection_function = detect_image_dual if use_dual_function else detect_image + + if file_path.suffix in [".jpg", ".jpeg", ".png"]: + try: + results = detection_function( + weight_path, + file_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + image = Image.open(file_path) + detected_image = self.draw_boxes(image, results) + detected_image.save("output.jpg") + self.imageLabel.setPixmap(QtGui.QPixmap("output.jpg")) + except Exception as e: + QtWidgets.QMessageBox.critical( + self, "Error", f"Error processing image: {str(e)}" + ) + elif file_path.suffix in [".zip"]: + try: + with zipfile.ZipFile(file_path, "r") as zip_ref: + zip_ref.extractall("temp_images") + detected_results = [] + csv_data = [] + image_files = list(Path("temp_images").glob("*")) + for img_path in tqdm(image_files, desc="Processing Images"): + if img_path.suffix in [".jpg", ".jpeg", ".png"]: + try: + results = detection_function( + weight_path, + img_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + image = Image.open(img_path) + detected_image = self.draw_boxes(image, results) + detected_results.append(detected_image) + for result in results: + box = result["bbox"] + confidence = result["conf"] + class_id = result["cls"] + csv_data.append( + [img_path.name, int(class_id), confidence, *box] + ) + except Exception as e: + print(f"Error processing image {img_path}: {str(e)}") + + csv_df = pd.DataFrame( + csv_data, + columns=["Filename", "Class", "Confidence", "X1", "Y1", "X2", "Y2"], + ) + save_path, _ = QFileDialog.getSaveFileName( + self, "Save CSV", "", "CSV Files (*.csv)" + ) + if save_path: + csv_df.to_csv(save_path, index=False) + else: + QtWidgets.QMessageBox.information( + self, "Cancelled", "CSV save cancelled" + ) + except Exception as e: + QtWidgets.QMessageBox.critical( + self, "Error", f"Error processing archive: {str(e)}" + ) + else: + QtWidgets.QMessageBox.warning(self, "Error", "Unsupported file format") + + +if __name__ == "__main__": + app = QApplication(sys.argv) + mainApp = MainApp() + mainApp.show() + sys.exit(app.exec_()) diff --git a/README.md b/README.md index 57fec4c7d..444971791 100644 --- a/README.md +++ b/README.md @@ -1,335 +1,39 @@ -# YOLOv9 +## Использование -Implementation of paper - [YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information](https://arxiv.org/abs/2402.13616) +1. Склонируйте репозиторий: -[![arxiv.org](http://img.shields.io/badge/cs.CV-arXiv%3A2402.13616-B31B1B.svg)](https://arxiv.org/abs/2402.13616) -[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/kadirnar/Yolov9) -[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/merve/yolov9) -[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov9-object-detection-on-custom-dataset.ipynb) -[![OpenCV](https://img.shields.io/badge/OpenCV-BlogPost-black?logo=opencv&labelColor=blue&color=black)](https://learnopencv.com/yolov9-advancing-the-yolo-legacy/) - -
- - - -
- - -## Performance - -MS COCO - -| Model | Test Size | APval | AP50val | AP75val | Param. | FLOPs | -| :-- | :-: | :-: | :-: | :-: | :-: | :-: | -| [**YOLOv9-T**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-t-converted.pt) | 640 | **38.3%** | **53.1%** | **41.3%** | **2.0M** | **7.7G** | -| [**YOLOv9-S**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-s-converted.pt) | 640 | **46.8%** | **63.4%** | **50.7%** | **7.1M** | **26.4G** | -| [**YOLOv9-M**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-m-converted.pt) | 640 | **51.4%** | **68.1%** | **56.1%** | **20.0M** | **76.3G** | -| [**YOLOv9-C**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c-converted.pt) | 640 | **53.0%** | **70.2%** | **57.8%** | **25.3M** | **102.1G** | -| [**YOLOv9-E**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e-converted.pt) | 640 | **55.6%** | **72.8%** | **60.6%** | **57.3M** | **189.0G** | - - - - -## Useful Links - -
Expand - -Custom training: https://github.com/WongKinYiu/yolov9/issues/30#issuecomment-1960955297 - -ONNX export: https://github.com/WongKinYiu/yolov9/issues/2#issuecomment-1960519506 https://github.com/WongKinYiu/yolov9/issues/40#issue-2150697688 https://github.com/WongKinYiu/yolov9/issues/130#issue-2162045461 - -ONNX export for segmentation: https://github.com/WongKinYiu/yolov9/issues/260#issue-2191162150 - -TensorRT inference: https://github.com/WongKinYiu/yolov9/issues/143#issuecomment-1975049660 https://github.com/WongKinYiu/yolov9/issues/34#issue-2150393690 https://github.com/WongKinYiu/yolov9/issues/79#issue-2153547004 https://github.com/WongKinYiu/yolov9/issues/143#issue-2164002309 - -QAT TensorRT: https://github.com/WongKinYiu/yolov9/issues/327#issue-2229284136 https://github.com/WongKinYiu/yolov9/issues/253#issue-2189520073 - -TensorRT inference for segmentation: https://github.com/WongKinYiu/yolov9/issues/446 - -TFLite: https://github.com/WongKinYiu/yolov9/issues/374#issuecomment-2065751706 - -OpenVINO: https://github.com/WongKinYiu/yolov9/issues/164#issue-2168540003 - -C# ONNX inference: https://github.com/WongKinYiu/yolov9/issues/95#issue-2155974619 - -C# OpenVINO inference: https://github.com/WongKinYiu/yolov9/issues/95#issuecomment-1968131244 - -OpenCV: https://github.com/WongKinYiu/yolov9/issues/113#issuecomment-1971327672 - -Hugging Face demo: https://github.com/WongKinYiu/yolov9/issues/45#issuecomment-1961496943 - -CoLab demo: https://github.com/WongKinYiu/yolov9/pull/18 - -ONNXSlim export: https://github.com/WongKinYiu/yolov9/pull/37 - -YOLOv9 ROS: https://github.com/WongKinYiu/yolov9/issues/144#issue-2164210644 - -YOLOv9 ROS TensorRT: https://github.com/WongKinYiu/yolov9/issues/145#issue-2164218595 - -YOLOv9 Julia: https://github.com/WongKinYiu/yolov9/issues/141#issuecomment-1973710107 - -YOLOv9 MLX: https://github.com/WongKinYiu/yolov9/issues/258#issue-2190586540 - -YOLOv9 StrongSORT with OSNet: https://github.com/WongKinYiu/yolov9/issues/299#issue-2212093340 - -YOLOv9 ByteTrack: https://github.com/WongKinYiu/yolov9/issues/78#issue-2153512879 - -YOLOv9 DeepSORT: https://github.com/WongKinYiu/yolov9/issues/98#issue-2156172319 - -YOLOv9 counting: https://github.com/WongKinYiu/yolov9/issues/84#issue-2153904804 - -YOLOv9 speed estimation: https://github.com/WongKinYiu/yolov9/issues/456 - -YOLOv9 face detection: https://github.com/WongKinYiu/yolov9/issues/121#issue-2160218766 - -YOLOv9 segmentation onnxruntime: https://github.com/WongKinYiu/yolov9/issues/151#issue-2165667350 - -Comet logging: https://github.com/WongKinYiu/yolov9/pull/110 - -MLflow logging: https://github.com/WongKinYiu/yolov9/pull/87 - -AnyLabeling tool: https://github.com/WongKinYiu/yolov9/issues/48#issue-2152139662 - -AX650N deploy: https://github.com/WongKinYiu/yolov9/issues/96#issue-2156115760 - -Conda environment: https://github.com/WongKinYiu/yolov9/pull/93 - -AutoDL docker environment: https://github.com/WongKinYiu/yolov9/issues/112#issue-2158203480 - -
- - -## Installation - -Docker environment (recommended) -
Expand - -``` shell -# create the docker container, you can change the share memory size if you have more. -nvidia-docker run --name yolov9 -it -v your_coco_path/:/coco/ -v your_code_path/:/yolov9 --shm-size=64g nvcr.io/nvidia/pytorch:21.11-py3 - -# apt install required packages -apt update -apt install -y zip htop screen libgl1-mesa-glx - -# pip install required packages -pip install seaborn thop - -# go to code folder -cd /yolov9 +```bash +docker build -t yolo9 . ``` -
- - -## Evaluation +3. Запустите контейнер с графическим процессором (монтируйте папку с данными): -[`yolov9-s-converted.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-s-converted.pt) [`yolov9-m-converted.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-m-converted.pt) [`yolov9-c-converted.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c-converted.pt) [`yolov9-e-converted.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e-converted.pt) -[`yolov9-s.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-s.pt) [`yolov9-m.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-m.pt) [`yolov9-c.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c.pt) [`yolov9-e.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt) -[`gelan-s.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-s.pt) [`gelan-m.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-m.pt) [`gelan-c.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c.pt) [`gelan-e.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-e.pt) - -``` shell -# evaluate converted yolov9 models -python val.py --data data/coco.yaml --img 640 --batch 32 --conf 0.001 --iou 0.7 --device 0 --weights './yolov9-c-converted.pt' --save-json --name yolov9_c_c_640_val - -# evaluate yolov9 models -# python val_dual.py --data data/coco.yaml --img 640 --batch 32 --conf 0.001 --iou 0.7 --device 0 --weights './yolov9-c.pt' --save-json --name yolov9_c_640_val - -# evaluate gelan models -# python val.py --data data/coco.yaml --img 640 --batch 32 --conf 0.001 --iou 0.7 --device 0 --weights './gelan-c.pt' --save-json --name gelan_c_640_val +```bash +docker run --gpus all -it -v /путь/к/вашей/папке:/workspace/mounted_folder yolo9 ``` -You will get the results: - -``` - Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.530 - Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.702 - Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.578 - Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.362 - Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.585 - Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.693 - Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.392 - Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.652 - Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.702 - Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.541 - Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.760 - Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.844 +Например: + ```bash +docker run --gpus all -it -v C:/Users/user/Desktop/data_and_weight:/workspace/mounted_folder yolo9 ``` -## Training - -Data preparation +4. Теперь, когда вы находитесь в контейнере, вы можете запустить `train.py` с помощью следующей команды, чтобы начать обучение модели: -``` shell -bash scripts/get_coco.sh +```bash +python train.py --batch 16 --epochs 300 --img 640 --device 0 --min-items 0 --data path/to/data.yaml --weights path/to/weights.pt --cfg models/detect/gelan-c.yaml --hyp hyp.scratch-high.yaml ``` -* Download MS COCO dataset images ([train](http://images.cocodataset.org/zips/train2017.zip), [val](http://images.cocodataset.org/zips/val2017.zip), [test](http://images.cocodataset.org/zips/test2017.zip)) and [labels](https://github.com/WongKinYiu/yolov7/releases/download/v0.1/coco2017labels-segments.zip). If you have previously used a different version of YOLO, we strongly recommend that you delete `train2017.cache` and `val2017.cache` files, and redownload [labels](https://github.com/WongKinYiu/yolov7/releases/download/v0.1/coco2017labels-segments.zip) - -Single GPU training - -``` shell -# train yolov9 models -python train_dual.py --workers 8 --device 0 --batch 16 --data data/coco.yaml --img 640 --cfg models/detect/yolov9-c.yaml --weights '' --name yolov9-c --hyp hyp.scratch-high.yaml --min-items 0 --epochs 500 --close-mosaic 15 - -# train gelan models -# python train.py --workers 8 --device 0 --batch 32 --data data/coco.yaml --img 640 --cfg models/detect/gelan-c.yaml --weights '' --name gelan-c --hyp hyp.scratch-high.yaml --min-items 0 --epochs 500 --close-mosaic 15 -``` - -Multiple GPU training - -``` shell -# train yolov9 models -python -m torch.distributed.launch --nproc_per_node 8 --master_port 9527 train_dual.py --workers 8 --device 0,1,2,3,4,5,6,7 --sync-bn --batch 128 --data data/coco.yaml --img 640 --cfg models/detect/yolov9-c.yaml --weights '' --name yolov9-c --hyp hyp.scratch-high.yaml --min-items 0 --epochs 500 --close-mosaic 15 - -# train gelan models -# python -m torch.distributed.launch --nproc_per_node 4 --master_port 9527 train.py --workers 8 --device 0,1,2,3 --sync-bn --batch 128 --data data/coco.yaml --img 640 --cfg models/detect/gelan-c.yaml --weights '' --name gelan-c --hyp hyp.scratch-high.yaml --min-items 0 --epochs 500 --close-mosaic 15 +Например: +```bash +python train.py --batch 32 --epochs 5 --img 640 --device 0 --min-items 0 --data mounted_folder/data/data.yaml --project mounted_folder/ --weights mounted_folder/gelan-c.pt --cfg models/detect/gelan-c.yaml --hyp hyp.scratch-high.yaml --workers 0 ``` +5. После окончания обучения в докере результат обучения сохраняется в папке "runs/train/exp". Чтобы перенести ее на вашу систему, выполните следующую команду (не в докере, а в терминале вашей системы) `docker cp`, указав путь к папке в контейнере и путь на вашей системе, куда вы хотите скопировать файл. -## Re-parameterization - -See [reparameterization.ipynb](https://github.com/WongKinYiu/yolov9/blob/main/tools/reparameterization.ipynb). - - -## Inference - -
- - - -
- -``` shell -# inference converted yolov9 models -python detect.py --source './data/images/horses.jpg' --img 640 --device 0 --weights './yolov9-c-converted.pt' --name yolov9_c_c_640_detect - -# inference yolov9 models -# python detect_dual.py --source './data/images/horses.jpg' --img 640 --device 0 --weights './yolov9-c.pt' --name yolov9_c_640_detect - -# inference gelan models -# python detect.py --source './data/images/horses.jpg' --img 640 --device 0 --weights './gelan-c.pt' --name gelan_c_c_640_detect +```bash +docker cp :workspace/runs/train/exp . ``` +Например: -## Citation - -``` -@article{wang2024yolov9, - title={{YOLOv9}: Learning What You Want to Learn Using Programmable Gradient Information}, - author={Wang, Chien-Yao and Liao, Hong-Yuan Mark}, - booktitle={arXiv preprint arXiv:2402.13616}, - year={2024} -} -``` - -``` -@article{chang2023yolor, - title={{YOLOR}-Based Multi-Task Learning}, - author={Chang, Hung-Shuo and Wang, Chien-Yao and Wang, Richard Robert and Chou, Gene and Liao, Hong-Yuan Mark}, - journal={arXiv preprint arXiv:2309.16921}, - year={2023} -} -``` - - -## Teaser - -Parts of code of [YOLOR-Based Multi-Task Learning](https://arxiv.org/abs/2309.16921) are released in the repository. - -
- - - -
- -#### Object Detection - -[`gelan-c-det.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c-det.pt) - -`object detection` - -``` shell -# coco/labels/{split}/*.txt -# bbox or polygon (1 instance 1 line) -python train.py --workers 8 --device 0 --batch 32 --data data/coco.yaml --img 640 --cfg models/detect/gelan-c.yaml --weights '' --name gelan-c-det --hyp hyp.scratch-high.yaml --min-items 0 --epochs 300 --close-mosaic 10 -``` - -| Model | Test Size | Param. | FLOPs | APbox | -| :-- | :-: | :-: | :-: | :-: | -| [**GELAN-C-DET**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c-det.pt) | 640 | 25.3M | 102.1G |**52.3%** | -| [**YOLOv9-C-DET**]() | 640 | 25.3M | 102.1G | **53.0%** | - -#### Instance Segmentation - -[`gelan-c-seg.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c-seg.pt) - -`object detection` `instance segmentation` - -``` shell -# coco/labels/{split}/*.txt -# polygon (1 instance 1 line) -python segment/train.py --workers 8 --device 0 --batch 32 --data coco.yaml --img 640 --cfg models/segment/gelan-c-seg.yaml --weights '' --name gelan-c-seg --hyp hyp.scratch-high.yaml --no-overlap --epochs 300 --close-mosaic 10 -``` - -| Model | Test Size | Param. | FLOPs | APbox | APmask | -| :-- | :-: | :-: | :-: | :-: | :-: | -| [**GELAN-C-SEG**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c-seg.pt) | 640 | 27.4M | 144.6G | **52.3%** | **42.4%** | -| [**YOLOv9-C-SEG**]() | 640 | 27.4M | 145.5G | **53.3%** | **43.5%** | - -#### Panoptic Segmentation - -[`gelan-c-pan.pt`](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c-pan.pt) - -`object detection` `instance segmentation` `semantic segmentation` `stuff segmentation` `panoptic segmentation` - -``` shell -# coco/labels/{split}/*.txt -# polygon (1 instance 1 line) -# coco/stuff/{split}/*.txt -# polygon (1 semantic 1 line) -python panoptic/train.py --workers 8 --device 0 --batch 32 --data coco.yaml --img 640 --cfg models/panoptic/gelan-c-pan.yaml --weights '' --name gelan-c-pan --hyp hyp.scratch-high.yaml --no-overlap --epochs 300 --close-mosaic 10 -``` - -| Model | Test Size | Param. | FLOPs | APbox | APmask | mIoU164k/10ksemantic | mIoUstuff | PQpanoptic | -| :-- | :-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: | -| [**GELAN-C-PAN**](https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c-pan.pt) | 640 | 27.6M | 146.7G | **52.6%** | **42.5%** | **39.0%/48.3%** | **52.7%** | **39.4%** | -| [**YOLOv9-C-PAN**]() | 640 | 28.8M | 187.0G | **52.7%** | **43.0%** | **39.8%/-** | **52.2%** | **40.5%** | - -#### Image Captioning (not yet released) - - - -`object detection` `instance segmentation` `semantic segmentation` `stuff segmentation` `panoptic segmentation` `image captioning` - -``` shell -# coco/labels/{split}/*.txt -# polygon (1 instance 1 line) -# coco/stuff/{split}/*.txt -# polygon (1 semantic 1 line) -# coco/annotations/*.json -# json (1 split 1 file) -python caption/train.py --workers 8 --device 0 --batch 32 --data coco.yaml --img 640 --cfg models/caption/gelan-c-cap.yaml --weights '' --name gelan-c-cap --hyp hyp.scratch-high.yaml --no-overlap --epochs 300 --close-mosaic 10 -``` - -| Model | Test Size | Param. | FLOPs | APbox | APmask | mIoU164k/10ksemantic | mIoUstuff | PQpanoptic | BLEU@4caption | CIDErcaption | -| :-- | :-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: | -| [**GELAN-C-CAP**]() | 640 | 47.5M | - | **51.9%** | **42.6%** | **42.5%/-** | **56.5%** | **41.7%** | **38.8** | **122.3** | -| [**YOLOv9-C-CAP**]() | 640 | 47.5M | - | **52.1%** | **42.6%** | **43.0%/-** | **56.4%** | **42.1%** | **39.1** | **122.0** | - - - -## Acknowledgements - -
Expand - -* [https://github.com/AlexeyAB/darknet](https://github.com/AlexeyAB/darknet) -* [https://github.com/WongKinYiu/yolor](https://github.com/WongKinYiu/yolor) -* [https://github.com/WongKinYiu/yolov7](https://github.com/WongKinYiu/yolov7) -* [https://github.com/VDIGPKU/DynamicDet](https://github.com/VDIGPKU/DynamicDet) -* [https://github.com/DingXiaoH/RepVGG](https://github.com/DingXiaoH/RepVGG) -* [https://github.com/ultralytics/yolov5](https://github.com/ultralytics/yolov5) -* [https://github.com/meituan/YOLOv6](https://github.com/meituan/YOLOv6) - -
diff --git a/create_dataset.py b/create_dataset.py new file mode 100644 index 000000000..ed421576a --- /dev/null +++ b/create_dataset.py @@ -0,0 +1,56 @@ +# Установите библиотеку Pillow, если она еще не установлена +# pip install pillow + +import os +import shutil +import glob + +def process_images(input_dir, weights_path, labels_dir): + # Создаем или очищаем директорию для изображений + image_dir = 'images' + if os.path.exists(image_dir): + shutil.rmtree(image_dir) + os.makedirs(image_dir) + + # Копируем файлы в директорию для изображений + for filename in os.listdir(input_dir): + input_path = os.path.join(input_dir, filename) + output_path = os.path.join(image_dir, filename) + shutil.copy(input_path, output_path) + + # Шаг 2: Запуск скрипта detect.py с указанием пути к весам и изображениям + os.system(f"python detect.py --weights {weights_path} --conf 0.6 --line-thickness 2 --source {image_dir} --device 0 --save-txt --save-conf --imgsz 640") + + # Шаг 3: Перемещение результатов в соответствующие папки + base_path = './runs/detect' # путь может быть другим в зависимости от настройки + exp_folders = glob.glob(os.path.join(base_path, 'exp*')) + exp_numbers = [int(folder.split('exp')[-1]) for folder in exp_folders if folder.split('exp')[-1].isdigit()] + max_exp = max(exp_numbers) if exp_numbers else None + + if max_exp is not None: + latest_folder = os.path.join(base_path, f'exp{max_exp}') + else: + latest_folder = os.path.join(base_path, 'exp') + + # Создаем директорию для лейблов, если она не существует + if not os.path.exists(labels_dir): + os.makedirs(labels_dir) + + # Перемещаем файлы лейблов в указанную папку labels + label_files = glob.glob(os.path.join(latest_folder, '*.txt')) + for label_file in label_files: + shutil.move(label_file, labels_dir) + + # Перемещаем файлы изображений в папку images + image_files = glob.glob(os.path.join(latest_folder, '*.jpg')) + for image_file in image_files: + shutil.move(image_file, image_dir) + + print(f"Всего изображений: {len(image_files)}") + print(f"Всего лейблов: {len(label_files)}") + +# Пример вызова функции +input_dir = r'C:\Users\user\Desktop\reports_orig' # укажите путь к вашей папке с изображениями +weights_path = r"C:\Users\user\Desktop\ddata\exp3\weights\best.pt" # укажите путь к вашему файлу весов +labels_dir = r'C:\Users\user\Desktop\labels' # укажите путь к папке, куда сохранять лейблы +process_images(input_dir, weights_path, labels_dir) diff --git a/del_labels.py b/del_labels.py new file mode 100644 index 000000000..3871c0df4 --- /dev/null +++ b/del_labels.py @@ -0,0 +1,35 @@ +import os + +def remove_confidence_from_labels(labels_dir): + # Проверяем, существует ли директория с лейблами + if not os.path.exists(labels_dir): + print(f"Директория {labels_dir} не существует.") + return + + # Получаем все файлы лейблов в директории + label_files = [f for f in os.listdir(labels_dir) if f.endswith('.txt')] + + for label_file in label_files: + label_path = os.path.join(labels_dir, label_file) + + # Читаем содержимое файла + with open(label_path, 'r') as f: + lines = f.readlines() + + # Убираем степень уверенности из каждого файла + new_lines = [] + for line in lines: + parts = line.strip().split() + if len(parts) == 6: + parts.pop(-1) # Удаляем последний элемент, который является степенью уверенности + new_lines.append(" ".join(parts) + "\n") + + # Сохраняем отредактированный файл + with open(label_path, 'w') as f: + f.writelines(new_lines) + + print(f"Обработано файлов лейблов: {len(label_files)}") + +# Пример вызова функции +labels_dir = r"C:\Users\user\Desktop\reports_detect_dataset\labels" # укажите путь к вашей папке с лейблами +remove_confidence_from_labels(labels_dir) diff --git a/del_no_label.py b/del_no_label.py new file mode 100644 index 000000000..eaa39781b --- /dev/null +++ b/del_no_label.py @@ -0,0 +1,33 @@ +import os + +def remove_images_without_labels(images_dir, labels_dir): + # Проверяем, существуют ли директории + if not os.path.exists(images_dir): + print(f"Директория {images_dir} не существует.") + return + + if not os.path.exists(labels_dir): + print(f"Директория {labels_dir} не существует.") + return + + # Получаем список всех файлов изображений и лейблов + image_files = [f for f in os.listdir(images_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] + label_files = [f for f in os.listdir(labels_dir) if f.endswith('.txt')] + + # Создаем множество базовых имен файлов лейблов (без расширения) + label_basenames = {os.path.splitext(f)[0] for f in label_files} + + # Удаляем изображения, для которых нет соответствующего файла лейбла + removed_count = 0 + for image_file in image_files: + image_basename = os.path.splitext(image_file)[0] + if image_basename not in label_basenames: + os.remove(os.path.join(images_dir, image_file)) + removed_count += 1 + + print(f"Удалено изображений: {removed_count}") + +# Пример вызова функции +images_dir = r"C:\Users\user\Desktop\reports_detect_dataset\images" # укажите путь к вашей папке с изображениями +labels_dir = r"C:\Users\user\Desktop\reports_detect_dataset\labels" # укажите путь к вашей папке с лейблами +remove_images_without_labels(images_dir, labels_dir) diff --git a/detect_function.py b/detect_function.py new file mode 100644 index 000000000..e2dba1ad1 --- /dev/null +++ b/detect_function.py @@ -0,0 +1,136 @@ +import os +import sys +from pathlib import Path + +import torch + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLO root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +from models.common import DetectMultiBackend +from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadCV2Image +from utils.general import Profile, check_img_size, non_max_suppression, scale_boxes +from utils.torch_utils import select_device, smart_inference_mode +from utils.augmentations import letterbox + + +import numpy as np +import cv2 + + +def sort_boxes_by_center(boxes): + def center(box): + xmin, ymin, xmax, ymax = box["bbox"] + center_x = (xmin + xmax) / 2 + center_y = (ymin + ymax) / 2 + return center_y, center_x + + return sorted(boxes, key=center) + + +@smart_inference_mode() +def run( + weights=Path("C:/Users/user/Desktop/yolov9/yolo.pt"), # model path or triton URL + source=None, # file/dir/URL/glob/screen/0(webcam) or numpy image + data=Path("C:/Users/user/Desktop/yolov9/data/coco.yaml"), # dataset.yaml path + imgsz=(640, 640), # inference size (height, width) + conf_thres=0.25, # confidence threshold + iou_thres=0.45, # NMS IOU threshold + max_det=1000, # maximum detections per image + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + classes=None, # filter by class: --class 0, or --class 0 2 3 + agnostic_nms=False, # class-agnostic NMS + augment=False, # augmented inference + visualize=False, # visualize features + half=False, # use FP16 half-precision inference + dnn=False, # use OpenCV DNN for ONNX inference, + sort_boxes=False, # sort boxes by center if True +): + device = select_device(device) + model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half) + stride, names, pt = model.stride, model.names, model.pt + imgsz = check_img_size(imgsz, s=stride) # check image size + + if isinstance(source, np.ndarray): + dataset = LoadCV2Image(source, img_size=imgsz, stride=stride, auto=True) + else: + source = str(source) + is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) + is_url = source.lower().startswith( + ("rtsp://", "rtmp://", "http://", "https://") + ) + webcam = ( + source.isnumeric() or source.endswith(".txt") or (is_url and not is_file) + ) + dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt) + + model.warmup(imgsz=(1 if pt or model.triton else 1, 3, *imgsz)) # warmup + results = [] + for path, im, im0s, vid_cap, s in dataset: + im = torch.from_numpy(im).to(model.device) + im = im.half() if model.fp16 else im.float() # uint8 to fp16/32 + im /= 255 # 0 - 255 to 0.0 - 1.0 + if len(im.shape) == 3: + im = im[None] # expand for batch dim + + pred = model(im, augment=augment, visualize=visualize) + pred = non_max_suppression( + pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det + ) + + for det in pred: # per image + if len(det): + det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0s.shape).round() + for *xyxy, conf, cls in reversed(det): + results.append( + { + "path": path, + "bbox": [int(x) for x in xyxy], + "conf": float(conf), + "cls": int(cls), + } + ) + + if sort_boxes: + results = sort_boxes_by_center(results) + + return results + + +def detect_image( + weights, + source, + data=None, + imgsz=(640, 640), + conf_thres=0.25, + iou_thres=0.45, + max_det=1000, + device="", + classes=None, + agnostic_nms=False, + augment=False, + visualize=False, + half=False, + dnn=False, + sort_boxes=False, +): + return run( + weights=weights, + source=source, + data=data, + imgsz=imgsz, + conf_thres=conf_thres, + iou_thres=iou_thres, + max_det=max_det, + device=device, + classes=classes, + agnostic_nms=agnostic_nms, + augment=augment, + visualize=visualize, + half=half, + dnn=dnn, + sort_boxes=sort_boxes, + ) diff --git a/detect_function_dual.py b/detect_function_dual.py new file mode 100644 index 000000000..8067a0ee7 --- /dev/null +++ b/detect_function_dual.py @@ -0,0 +1,143 @@ +import os +import sys +from pathlib import Path + +import torch + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLO root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +from models.common import DetectMultiBackend +from utils.dataloaders import ( + IMG_FORMATS, + VID_FORMATS, + LoadImages, + LoadScreenshots, + LoadStreams, +) +from utils.general import ( + check_file, + check_img_size, + non_max_suppression, + scale_boxes, + strip_optimizer, + xyxy2xywh, + LOGGER, +) +from utils.plots import Annotator, colors, save_one_box +from utils.torch_utils import select_device, smart_inference_mode + + +@smart_inference_mode() +def run_detect( + weights, # model path(s) + source, # file/dir/URL/glob/screen/0(webcam) + data=None, # dataset.yaml path, can be None for detection only + imgsz=640, # inference size (pixels) + conf_thres=0.25, # confidence threshold + iou_thres=0.45, # NMS IoU threshold + max_det=1000, # maximum detections per image + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + half=False, # use FP16 half-precision inference + dnn=False, # use OpenCV DNN for ONNX inference + sort_boxes=False, # sort boxes by center if True +): + source = str(source) + is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) + is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://")) + webcam = source.isnumeric() or source.endswith(".txt") or (is_url and not is_file) + screenshot = source.lower().startswith("screen") + if is_url and is_file: + source = check_file(source) # download + + # Load model + device = select_device(device) + model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half) + stride, names, pt = model.stride, model.names, model.pt + imgsz = check_img_size(imgsz, s=stride) # check image size + + # Dataloader + if webcam: + dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt) + elif screenshot: + dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) + else: + dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt) + + # Run inference + model.warmup(imgsz=(1 if pt else len(dataset), 3, *imgsz)) # warmup + results = [] + for path, im, im0s, vid_cap, s in dataset: + im = torch.from_numpy(im).to(model.device) + im = im.half() if model.fp16 else im.float() # uint8 to fp16/32 + im /= 255 # 0 - 255 to 0.0 - 1.0 + if len(im.shape) == 3: + im = im[None] # expand for batch dim + + # Inference + pred = model(im, augment=False, visualize=False) + pred = non_max_suppression(pred[0][1], conf_thres, iou_thres, max_det=max_det) + + # Process predictions + for i, det in enumerate(pred): # per image + p, im0 = path, im0s + if webcam: # batch_size >= 1 + p, im0 = path[i], im0s[i] + + p = Path(p) # to Path + gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh + imc = im0.copy() # for save_crop + + if len(det): + det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round() + for *xyxy, conf, cls in reversed(det): + results.append( + { + "path": str(p), + "bbox": [int(x) for x in xyxy], + "conf": float(conf), + "cls": int(cls), + } + ) + + if sort_boxes: + results = sorted( + results, + key=lambda x: ( + (x["bbox"][0] + x["bbox"][2]) / 2, + (x["bbox"][1] + x["bbox"][3]) / 2, + ), + ) + + return results + + +def detect_image_dual( + weights, + source, + data=None, # make data optional + imgsz=(640, 640), + conf_thres=0.25, + iou_thres=0.45, + max_det=1000, + device="", + half=False, + dnn=False, + sort_boxes=False, +): + return run_detect( + weights=weights, + source=source, + data=data, + imgsz=imgsz, + conf_thres=conf_thres, + iou_thres=iou_thres, + max_det=max_det, + device=device, + half=half, + dnn=dnn, + sort_boxes=sort_boxes, + ) diff --git a/draw_boxes.py b/draw_boxes.py new file mode 100644 index 000000000..aaf195abd --- /dev/null +++ b/draw_boxes.py @@ -0,0 +1,148 @@ +import os +from pathlib import Path +from PIL import Image, ImageDraw, ImageFont, ImageColor + +# Путь к весам модели +weight_path = r"C:\Users\user\Desktop\crispi_defects\data\exp3\weights\best.pt" + +# Заданные в коде названия классов +class_names = ["Пузырь", "Трещина"] + +# Расширенный список цветов для каждого класса с альфа-каналом (прозрачность) +default_colors = [ + (0, 0, 0, 128), # Чёрный с прозрачностью + (0, 0, 139, 128), # Тёмно-синий с прозрачностью + (139, 0, 0, 128), # Тёмно-красный с прозрачностью + (0, 100, 0, 128), # Тёмно-зелёный с прозрачностью + (75, 0, 130, 128), # Индиго с прозрачностью + (47, 79, 79, 128), # Тёмный серо-зелёный с прозрачностью + (139, 0, 139, 128), # Тёмная магента с прозрачностью + (139, 69, 19, 128), # Седло-коричневый с прозрачностью + (128, 0, 0, 128), # Бордовый с прозрачностью + (72, 61, 139, 128), # Тёмный серо-синий с прозрачностью + (70, 130, 180, 128), # Стальной синий с прозрачностью + (85, 107, 47, 128), # Тёмный оливковый с прозрачностью + (106, 90, 205, 128), # Сланцевый синий с прозрачностью + (46, 139, 87, 128), # Морская волна с прозрачностью + (160, 82, 45, 128), # Сиенна с прозрачностью +] + +# Функция для рисования легенды на изображении +def draw_legend(image, class_names, colors, font): + draw = ImageDraw.Draw(image, "RGBA") + padding = 10 + + # Вычисляем высоту легенды и ширину на основе размеров текста + legend_height = len(class_names) * (font.getbbox(class_names[0])[3] + padding) + padding + legend_width = max(font.getbbox(name)[2] - font.getbbox(name)[0] for name in class_names) + padding * 4 + 40 + + # Координаты верхнего правого угла + x1 = image.width - legend_width - padding + y1 = padding + x2 = image.width - padding + y2 = y1 + legend_height + + # Рисуем фон для легенды + draw.rectangle([x1, y1, x2, y2], fill=(255, 255, 255, 200)) + + for i, (name, color) in enumerate(zip(class_names, colors)): + text_position = (x1 + padding * 2 + 30, y1 + padding + i * (font.getbbox(name)[3] + padding)) + color_position = (x1 + padding, text_position[1] + (font.getbbox(name)[3] - font.getbbox(name)[1] - 20) // 2) + + # Рисуем цветной квадрат + draw.rectangle([color_position, (color_position[0] + 20, color_position[1] + 20)], fill=color) + + # Рисуем текст + draw.text(text_position, name, fill=(0, 0, 0, 255), font=font) + + return image + +# Функция для рисования истинных боксов на изображении +def draw_true_boxes_on_image(image, true_boxes, start_index=1): + draw = ImageDraw.Draw(image, "RGBA") + font = ImageFont.truetype("arial.ttf", size=16) + width, height = image.size + + # Рисуем истинные боксы (предполагается, что координаты нормализованы) + for i, box in enumerate(true_boxes, start=start_index): + class_id = int(box["cls"]) + + # Преобразуем относительные координаты в абсолютные + x_center, y_center, bbox_width, bbox_height = box["bbox"] + x1 = int((x_center - bbox_width / 2) * width) + y1 = int((y_center - bbox_height / 2) * height) + x2 = int((x_center + bbox_width / 2) * width) + y2 = int((y_center + bbox_height / 2) * height) + + # Проверка, что боксы попадают в видимую область изображения + if x1 < 0 or y1 < 0 or x2 > width or y2 > height: + print(f"Warning: True box for class {class_names[class_id]} is out of image bounds: ({x1}, {y1}), ({x2}, {y2})") + continue + + # Используем цвет, соответствующий классу, с прозрачностью + color = default_colors[class_id % len(default_colors)] + + # Рисуем более тонкий полупрозрачный прямоугольник + draw.rectangle([x1, y1, x2, y2], outline=color, width=2) + label = f"{i}" # Только номер + + # Центрирование текста по высоте бокса и размещение его слева + text_bbox = draw.textbbox((0, 0), label, font=font) + text_width = text_bbox[2] - text_bbox[0] + text_height = text_bbox[3] - text_bbox[1] + + text_x = x1 - text_width - 5 # Размещаем текст левее бокса с отступом + text_y = y1 + (bbox_height * height - text_height) / 2 # Центрируем текст по высоте бокса + + # Добавление прозрачного фона под текст для лучшей видимости + text_bg_color = (color[0], color[1], color[2], 128) # Прозрачный фон под текст + draw.rectangle([text_x - 2, text_y - 2, text_x + text_width + 2, text_y + text_height + 2], fill=text_bg_color) + draw.text((text_x, text_y), label, fill=(255, 255, 255, 128), font=font) # Прозрачный белый текст + + # Добавляем легенду на изображение + image = draw_legend(image, class_names, default_colors, font) + + return image, len(true_boxes) + +# Функция обработки изображений из указанной папки +def process_images_in_folder(input_folder, output_folder, start_index=1): + image_folder = Path(input_folder) / "images" + label_folder = Path(input_folder) / "labels" + output_combined_folder = Path(output_folder) / "combined_boxes" + + # Создаем выходные папки, если их нет + output_combined_folder.mkdir(parents=True, exist_ok=True) + + current_index = start_index + + # Обрабатываем все изображения в папке + for image_path in image_folder.glob("*.*"): + if image_path.suffix.lower() not in [".jpg", ".jpeg", ".png"]: + continue + + image = Image.open(image_path).convert("RGB") + label_path = label_folder / (image_path.stem + ".txt") + + # Чтение истинных боксов (если необходимо) + true_boxes = [] + if label_path.exists(): + with open(label_path, 'r') as f: + for line in f: + parts = line.strip().split() + class_id = int(parts[0]) + bbox = list(map(float, parts[1:])) + true_boxes.append({"cls": class_id, "bbox": bbox}) + + if true_boxes: + # Рисуем истинные боксы на изображении и обновляем текущий индекс + combined_image, num_boxes = draw_true_boxes_on_image(image, true_boxes, start_index=current_index) + combined_image.save(output_combined_folder / image_path.name) + current_index += num_boxes + + print("Обработка завершена. Последний индекс бокса:", current_index - 1) + +# Использование функции +input_folder = r"C:\Users\user\Desktop\rusal_data\data\test" # Путь до папки с папками images и labels +output_folder = r"C:\Users\user\Desktop\rusal_data\out" # Путь для сохранения изображений с боксами + +process_images_in_folder(input_folder, output_folder) diff --git a/finc_clone.py b/finc_clone.py new file mode 100644 index 000000000..c198049e4 --- /dev/null +++ b/finc_clone.py @@ -0,0 +1,45 @@ +import os +import hashlib + +def get_file_hash(file_path): + """Возвращает хеш содержимого файла.""" + hasher = hashlib.md5() + with open(file_path, 'rb') as f: + buf = f.read() + hasher.update(buf) + return hasher.hexdigest() + +def check_duplicate_labels(labels_dir): + # Проверяем, существует ли директория + if not os.path.exists(labels_dir): + print(f"Директория {labels_dir} не существует.") + return + + # Получаем список всех файлов лейблов + label_files = [f for f in os.listdir(labels_dir) if f.endswith('.txt')] + + # Создаем словарь для хранения хешей файлов + hash_dict = {} + + duplicates = [] + + for label_file in label_files: + label_path = os.path.join(labels_dir, label_file) + file_hash = get_file_hash(label_path) + + if file_hash in hash_dict: + duplicates.append((hash_dict[file_hash], label_file)) + else: + hash_dict[file_hash] = label_file + + if duplicates: + print("Найдены файлы с одинаковой разметкой:") + for original, duplicate in duplicates: + print(f"{duplicate} является дубликатом {original}") + else: + print("Дубликаты не найдены.") + +# Пример вызова функции +labels_dir = r"C:\Users\user\Desktop\reports_detect_dataset — копия\labels" + # укажите путь к вашей папке с лейблами +check_duplicate_labels(labels_dir) \ No newline at end of file diff --git a/gradio_main.py b/gradio_main.py new file mode 100644 index 000000000..c8eb0dfdd --- /dev/null +++ b/gradio_main.py @@ -0,0 +1,180 @@ +import gradio as gr +import zipfile +from pathlib import Path +from PIL import Image, ImageDraw, ImageFont +import tkinter as tk +from tkinter import filedialog +import pandas as pd +from tqdm import tqdm +from detect_function_dual import detect_image_dual +from detect_function import detect_image + +# Переключатель для выбора функции детекции +use_dual_function = ( + True # Установите в False для использования стандартной функции детекции +) + +# Путь к весам модели +weight_path = ( + r"C:\Users\user\Desktop\crispi_defects\data\exp3\weights\best.pt" +) + + +# Функция для рисования боксов на изображении +def draw_boxes(image, results): + draw = ImageDraw.Draw(image) + font = ImageFont.load_default() + for result in results: + x1, y1, x2, y2 = result["bbox"] + confidence = result["conf"] + class_id = result["cls"] + draw.rectangle([x1, y1, x2, y2], outline="red", width=2) + text = f"Class {int(class_id)} {confidence:.2f}" + text_size = draw.textbbox((0, 0), text, font=font) + text_location = [x1, y1 - text_size[3]] + if text_location[1] < 0: + text_location[1] = y1 + text_size[3] + draw.rectangle([x1, y1 - text_size[3], x1 + text_size[2], y1], fill="red") + draw.text((x1, y1 - text_size[3]), text, fill="white", font=font) + return image + + +# Функция для обработки одного файла +def process_file(file, conf_thres, iou_thres): + if file is None or not file.name: + print("Файл не предоставлен") + return None + + try: + file_path = Path(file.name) + detection_function = detect_image_dual if use_dual_function else detect_image + + print( + f"Используется функция: {'detect_image_dual' if use_dual_function else 'detect_image'}" + ) + print(f"Параметры: conf_thres={conf_thres}, iou_thres={iou_thres}") + print(f"Обрабатываемый файл: {file_path}") + + if file_path.suffix in [".jpg", ".jpeg", ".png"]: + try: + print(f"Обработка изображения: {file_path}") + results = detection_function( + weight_path, + file_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + print(f"Результаты детекции: {results}") + image = Image.open(file) + detected_image = draw_boxes(image, results) + return detected_image + except Exception as e: + print(f"Произошла ошибка при обработке изображения: {str(e)}") + return None + elif file_path.suffix in [".zip"]: + try: + with zipfile.ZipFile(file, "r") as zip_ref: + zip_ref.extractall("temp_images") + detected_results = [] + csv_data = [] + image_files = list(Path("temp_images").glob("*")) + for img_path in tqdm(image_files, desc="Обработка изображений"): + if img_path.suffix in [".jpg", ".jpeg", ".png"]: + try: + print(f"Обработка изображения из архива: {img_path}") + results = detection_function( + weight_path, + img_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + print( + f"Результаты детекции для {img_path}: {results.keys()}" + ) + image = Image.open(img_path) + detected_image = draw_boxes(image, results) + detected_results.append(detected_image) + # Добавляем информацию о боксах в CSV данные + for result in results: + box = result["bbox"] + confidence = result["conf"] + class_id = result["cls"] + csv_data.append( + [ + img_path.name, + int(class_id), + confidence, + *box, + ] + ) + except Exception as e: + print( + f"Произошла ошибка при обработке изображения {img_path}: {str(e)}" + ) + + csv_df = pd.DataFrame( + csv_data, + columns=["Filename", "Class", "Confidence", "X1", "Y1", "X2", "Y2"], + ) + + # Открытие окна выбора пути для сохранения CSV файла + root = tk.Tk() + root.withdraw() + root.lift() # Поднятие окна поверх других + root.attributes("-topmost", True) + default_filename = f"{file_path.stem}_results.csv" + save_path = filedialog.asksaveasfilename( + initialfile=default_filename, + defaultextension=".csv", + filetypes=[("CSV files", "*.csv")], + ) + + if save_path: + print(f"Сохранение CSV файла в: {save_path}") + csv_df.to_csv(save_path, index=False) + return None + else: + print("Сохранение CSV файла отменено") + return None + except Exception as e: + print(f"Произошла ошибка при обработке архива: {str(e)}") + return None + else: + print("Неподдерживаемый формат файла.") + return None + except Exception as e: + print(f"Произошла общая ошибка: {str(e)}") + return None + + +def reset_interface(): + return None + + +# Gradio интерфейс +with gr.Blocks() as demo: + with gr.Row(): + file_input = gr.File( + label="Загрузите изображение или архив", file_count="single" + ) + with gr.Column(): + conf_thres_input = gr.Slider(0, 1, value=0.25, label="Confidence Threshold") + iou_thres_input = gr.Slider(0, 1, value=0.45, label="IoU Threshold") + + output_image = gr.Image(label="Детектированное изображение") + + def process_file_wrapper(file, conf_thres, iou_thres): + return process_file(file, conf_thres, iou_thres) + + file_input.change( + process_file_wrapper, + inputs=[file_input, conf_thres_input, iou_thres_input], + outputs=[output_image], + ) + + file_input.clear(fn=reset_interface, inputs=None, outputs=[output_image]) + +# Запуск интерфейса с параметром share=False +demo.launch(share=True) diff --git a/gradio_main2.py b/gradio_main2.py new file mode 100644 index 000000000..c388849a3 --- /dev/null +++ b/gradio_main2.py @@ -0,0 +1,122 @@ +import gradio as gr +from pathlib import Path +from PIL import Image, ImageDraw, ImageFont +from tqdm import tqdm +from detect_function_dual import detect_image_dual +from detect_function import detect_image + +# Переключатель для выбора функции детекции +use_dual_function = True # Установите в False для использования стандартной функции детекции + +# Путь к весам модели +weight_path = r"C:\Users\user\Desktop\crispi_defects\data\exp3\weights\best.pt" + +# Заданные в коде названия классов +class_names = ["Вмятина", "На проволоке", "Накол"] + +# Расширенный список цветов для каждого класса +default_colors = [ + "#000000", # Чёрный + "#00008B", # Тёмно-синий + "#8B0000", # Тёмно-красный + "#006400", # Тёмно-зелёный + "#4B0082", # Индиго + "#2F4F4F", # Тёмный серо-зелёный + "#8B008B", # Тёмная магента + "#8B4513", # Седло-коричневый + "#800000", # Бордовый + "#483D8B", # Тёмный серо-синий + "#4682B4", # Стальной синий + "#556B2F", # Тёмный оливковый + "#6A5ACD", # Сланцевый синий + "#2E8B57", # Морская волна + "#A0522D", # Сиенна +] + +# Функция для рисования боксов на изображении +def draw_boxes(image, results): + draw = ImageDraw.Draw(image) + font = ImageFont.truetype("arial.ttf", size=16) + for result in results: + x1, y1, x2, y2 = result["bbox"] + confidence = result["conf"] + class_id = int(result["cls"]) + class_name = class_names[class_id] if class_id < len(class_names) else f"Class {class_id}" + color = default_colors[class_id % len(default_colors)] + + # Рисуем прямоугольник + draw.rectangle([x1, y1, x2, y2], outline=color, width=3) + + # Подготавливаем текст метки + label = f"{class_name}: {confidence:.2f}" + text_bbox = draw.textbbox((0, 0), label, font=font) + text_width = text_bbox[2] - text_bbox[0] + text_height = text_bbox[3] - text_bbox[1] + text_background = [x1, y1 - text_height - 4, x1 + text_width + 4, y1] + + # Рисуем фон для текста + draw.rectangle(text_background, fill=color) + + # Рисуем текст метки + draw.text((x1 + 2, y1 - text_height - 2), label, fill="white", font=font) + + return image + +# Функция для обработки нескольких файлов +def process_files(files, gallery_state): + if not files: + return gallery_state, "Файлы не предоставлены." + + detection_function = detect_image_dual if use_dual_function else detect_image + + # Устанавливаем пороги внутри кода + conf_thres = 0.4 + iou_thres = 0.4 + + for file_path in tqdm(files, desc="Обработка изображений"): + file_path = Path(file_path) # Преобразуем каждый файл в объект Path + print(f"Обрабатываемый файл: {file_path}") + + if file_path.suffix.lower() in [".jpg", ".jpeg", ".png"]: + image = Image.open(file_path).convert("RGB") + results = detection_function( + weight_path, + file_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + detected_image = draw_boxes(image, results) + gallery_state.append(detected_image) + + return gallery_state, None + +def reset_interface(): + return [], None + +# Gradio интерфейс +with gr.Blocks() as demo: + gr.Markdown("

Интерфейс обнаружения объектов

") + + with gr.Row(): + with gr.Column(): + file_input = gr.File( + label="Загрузите изображения", + file_count="multiple", # Разрешаем загружать несколько файлов + type="filepath", # Используем 'filepath', чтобы получить пути к файлам + ) + submit_button = gr.Button("Обработать") + + with gr.Column(): + gallery_state = gr.State([]) # Для хранения изображений + output_image = gr.Gallery(label="Галерея изображений") + + submit_button.click( + process_files, + inputs=[file_input, gallery_state], + outputs=[output_image, file_input], + ) + + demo.load(fn=reset_interface, outputs=[gallery_state, file_input]) + +demo.launch(share=True) diff --git a/gradio_main_old.py b/gradio_main_old.py new file mode 100644 index 000000000..cbc26a6d2 --- /dev/null +++ b/gradio_main_old.py @@ -0,0 +1,132 @@ +import gradio as gr +import zipfile +import os +from pathlib import Path +from PIL import Image, ImageDraw, ImageFont +import yolov9 +import tkinter as tk +from tkinter import filedialog +import numpy as np +import pandas as pd +from tqdm import tqdm + +# Загрузка модели YOLOv9 +model = yolov9.load( + r"C:\Users\pasha\OneDrive\Рабочий стол\yolo_weights\yolo_word_detectino21.pt", + device="cpu", +) +model.conf = 0.25 # Порог уверенности NMS +model.iou = 0.45 # Порог IoU NMS + + +def draw_boxes(image, boxes, confidences, class_ids, class_names): + draw = ImageDraw.Draw(image) + font = ImageFont.load_default() + for box, confidence, class_id in zip(boxes, confidences, class_ids): + x1, y1, x2, y2 = box + draw.rectangle([x1, y1, x2, y2], outline="red", width=2) + text = f"{class_names[int(class_id)]} {confidence:.2f}" + text_size = draw.textbbox((0, 0), text, font=font) + text_location = [x1, y1 - text_size[3]] + if text_location[1] < 0: + text_location[1] = y1 + text_size[3] + draw.rectangle([x1, y1 - text_size[3], x1 + text_size[2], y1], fill="red") + draw.text((x1, y1 - text_size[3]), text, fill="white", font=font) + return image + + +def process_file(file): + if file is None: + return None + + file_path = Path(file.name) + class_names = model.names # Получаем названия классов + if file_path.suffix in [".jpg", ".jpeg", ".png"]: + # Если файл - изображение + image = Image.open(file) + results = model(np.array(image)) + boxes = results.pred[0][:, :4].cpu().numpy() # Получаем координаты боксов + confidences = results.pred[0][:, 4].cpu().numpy() # Получаем уверенности + class_ids = ( + results.pred[0][:, 5].cpu().numpy() + ) # Получаем идентификаторы классов + detected_image = draw_boxes(image, boxes, confidences, class_ids, class_names) + return detected_image + elif file_path.suffix in [".zip"]: + try: + with zipfile.ZipFile(file, "r") as zip_ref: + zip_ref.extractall("temp_images") + detected_results = [] + csv_data = [] + image_files = list(Path("temp_images").glob("*")) + for img_path in tqdm(image_files, desc="Обработка изображений"): + if img_path.suffix in [".jpg", ".jpeg", ".png"]: + image = Image.open(img_path) + results = model(np.array(image)) + boxes = ( + results.pred[0][:, :4].cpu().numpy() + ) # Получаем координаты боксов + confidences = ( + results.pred[0][:, 4].cpu().numpy() + ) # Получаем уверенности + class_ids = ( + results.pred[0][:, 5].cpu().numpy() + ) # Получаем идентификаторы классов + detected_image = draw_boxes( + image, boxes, confidences, class_ids, class_names + ) + detected_results.append(detected_image) + # Добавляем информацию о боксах в CSV данные + for box, confidence, class_id in zip(boxes, confidences, class_ids): + csv_data.append( + [ + img_path.name, + class_names[int(class_id)], + confidence, + *box, + ] + ) + + csv_df = pd.DataFrame( + csv_data, + columns=["Filename", "Class", "Confidence", "X1", "Y1", "X2", "Y2"], + ) + + # Открытие окна выбора пути для сохранения CSV файла + root = tk.Tk() + root.withdraw() + root.lift() # Поднятие окна поверх других + root.attributes("-topmost", True) + default_filename = f"{file_path.stem}_results.csv" + save_path = filedialog.asksaveasfilename( + initialfile=default_filename, + defaultextension=".csv", + filetypes=[("CSV files", "*.csv")], + ) + + if save_path: + csv_df.to_csv(save_path, index=False) + return None + else: + return None + except Exception as e: + print(f"Произошла ошибка при обработке архива: {str(e)}") + return None + else: + print("Неподдерживаемый формат файла.") + return None + + +# Gradio интерфейс +with gr.Blocks() as demo: + with gr.Row(): + file_input = gr.File( + label="Загрузите изображение или архив", file_count="single" + ) + + output_image = gr.Image(label="Детектированное изображение") + + file_input.change(process_file, inputs=file_input, outputs=[output_image]) + +# Запуск интерфейса с параметром share=True +demo.launch(share=False) diff --git a/output.jpg b/output.jpg new file mode 100644 index 000000000..70dd2dfe5 Binary files /dev/null and b/output.jpg differ diff --git a/predct_images.py b/predct_images.py new file mode 100644 index 000000000..55b80ba90 --- /dev/null +++ b/predct_images.py @@ -0,0 +1,136 @@ +import os +from pathlib import Path +from PIL import Image, ImageDraw, ImageFont +import shutil +from detect_function_dual import detect_image_dual +from detect_function import detect_image + +# Путь к весам модели +weight_path = r"C:\Users\user\Desktop\crispi_defects\data\exp3\weights\best.pt" + +# Заданные в коде названия классов +class_names = ["Вмятина", "На проволоке", "Накол"] + +# Расширенный список цветов для каждого класса +default_colors = [ + "#000000", # Чёрный + "#00008B", # Тёмно-синий + "#8B0000", # Тёмно-красный + "#006400", # Тёмно-зелёный + "#4B0082", # Индиго + "#2F4F4F", # Тёмный серо-зелёный + "#8B008B", # Тёмная магента + "#8B4513", # Седло-коричневый + "#800000", # Бордовый + "#483D8B", # Тёмный серо-синий + "#4682B4", # Стальной синий + "#556B2F", # Тёмный оливковый + "#6A5ACD", # Сланцевый синий + "#2E8B57", # Морская волна + "#A0522D", # Сиенна +] + +# Функция для рисования легенды на изображении +def draw_legend(image, class_names, colors, font): + draw = ImageDraw.Draw(image, "RGBA") + padding = 10 + + # Вычисляем высоту легенды и ширину на основе размеров текста + legend_height = len(class_names) * (font.getbbox(class_names[0])[3] + padding) + padding + legend_width = max(font.getbbox(name)[2] - font.getbbox(name)[0] for name in class_names) + padding * 4 + 40 + + # Координаты верхнего правого угла + x1 = image.width - legend_width - padding + y1 = padding + x2 = image.width - padding + y2 = y1 + legend_height + + # Рисуем фон для легенды + draw.rectangle([x1, y1, x2, y2], fill=(255, 255, 255, 200)) + + for i, (name, color) in enumerate(zip(class_names, colors)): + text_position = (x1 + padding * 2 + 30, y1 + padding + i * (font.getbbox(name)[3] + padding)) + color_position = (x1 + padding, text_position[1] + (font.getbbox(name)[3] - font.getbbox(name)[1] - 20) // 2) + + # Рисуем цветной квадрат + draw.rectangle([color_position, (color_position[0] + 20, color_position[1] + 20)], fill=color) + + # Рисуем текст + draw.text(text_position, name, fill=(0, 0, 0, 255), font=font) + + return image + +# Функция для рисования предсказанных боксов на изображении +def draw_boxes_on_image(image, pred_boxes): + draw = ImageDraw.Draw(image) + font = ImageFont.truetype("arial.ttf", size=16) + width, height = image.size + + # Рисуем предсказанные боксы (предполагается, что координаты уже абсолютные) + for box in pred_boxes: + class_id = int(box["cls"]) + confidence = box["conf"] + + # Используем абсолютные координаты напрямую + x1, y1, x2, y2 = map(int, box["bbox"]) + + # Проверка, что боксы попадают в видимую область изображения + if x1 < 0 or y1 < 0 or x2 > width or y2 > height: + print(f"Warning: Predicted box for class {class_names[class_id]} is out of image bounds: ({x1}, {y1}), ({x2}, {y2})") + continue + + # Рисуем прямоугольник для предсказанного бокса + color = default_colors[class_id % len(default_colors)] + draw.rectangle([x1, y1, x2, y2], outline=color, width=3) + + # Подготавливаем текст с именем класса и вероятностью + label = f"{class_names[class_id]}: {confidence:.2f}" + text_bbox = draw.textbbox((x1, y1 - 10), label, font=font) + draw.rectangle([text_bbox[0] - 2, text_bbox[1] - 2, text_bbox[2] + 2, text_bbox[3] + 2], fill=color) + draw.text((x1, y1 - 10), label, fill="white", font=font) + + # Добавляем легенду на изображение + image = draw_legend(image, class_names, default_colors, font) + + return image + + +# Функция обработки изображений из указанной папки +def process_images_in_folder(input_folder, output_folder, conf_thres=0.4, iou_thres=0.45, use_dual_function=True): + image_folder = Path(input_folder) / "images" + output_combined_folder = Path(output_folder) / "combined_boxes" + + # Создаем выходные папки, если их нет + output_combined_folder.mkdir(parents=True, exist_ok=True) + + # Выбор функции детекции + detection_function = detect_image_dual if use_dual_function else detect_image + + # Обрабатываем все изображения в папке + for image_path in image_folder.glob("*.*"): + if image_path.suffix.lower() not in [".jpg", ".jpeg", ".png"]: + continue + + image = Image.open(image_path).convert("RGB") + + # Получение предсказанных боксов + pred_boxes = detection_function( + weight_path, + image_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + + # Рисуем предсказанные боксы на изображении + combined_image = image.copy() + combined_image = draw_boxes_on_image(combined_image, pred_boxes) + combined_image.save(output_combined_folder / image_path.name) + + print("Обработка завершена.") + +# Использование функции +input_folder = r"C:\Users\user\Desktop\crispi_defects\data\test" # Путь до папки с папками images и labels +output_folder = r"C:\Users\user\Desktop\Проволока\снимки для КРИТБИ\yoloimages\out_test" # Путь для сохранения изображений с боксами + +process_images_in_folder(input_folder, output_folder) diff --git a/predict_images_true_not.py b/predict_images_true_not.py new file mode 100644 index 000000000..7879cb86c --- /dev/null +++ b/predict_images_true_not.py @@ -0,0 +1,146 @@ +import os +from pathlib import Path +from PIL import Image, ImageDraw, ImageFont +import shutil +from detect_function_dual import detect_image_dual +from detect_function import detect_image + +# Путь к весам модели +weight_path = r"C:\Users\user\Desktop\crispi_defects\data\exp3\weights\best.pt" + +# Заданные в коде названия классов +class_names = ["Вмятина", "На проволоке", "Накол"] + +# Расширенный список цветов для каждого класса +default_colors = [ + "#000000", # Чёрный + "#00008B", # Тёмно-синий + "#8B0000", # Тёмно-красный + "#006400", # Тёмно-зелёный + "#4B0082", # Индиго + "#2F4F4F", # Тёмный серо-зелёный + "#8B008B", # Тёмная магента + "#8B4513", # Седло-коричневый + "#800000", # Бордовый + "#483D8B", # Тёмный серо-синий + "#4682B4", # Стальной синий + "#556B2F", # Тёмный оливковый + "#6A5ACD", # Сланцевый синий + "#2E8B57", # Морская волна + "#A0522D", # Сиенна +] + +# Функция для рисования боксов на изображении +def draw_boxes_on_image(image, true_boxes, pred_boxes, true_color="#00FF00", pred_color="#FF0000"): + draw = ImageDraw.Draw(image) + font = ImageFont.truetype("arial.ttf", size=16) + width, height = image.size + + # Рисуем истинные боксы (предполагается, что координаты нормализованы) + for box in true_boxes: + class_id = int(box["cls"]) + + # Преобразуем относительные координаты в абсолютные + x_center, y_center, bbox_width, bbox_height = box["bbox"] + x1 = int((x_center - bbox_width / 2) * width) + y1 = int((y_center - bbox_height / 2) * height) + x2 = int((x_center + bbox_width / 2) * width) + y2 = int((y_center + bbox_height / 2) * height) + + # Проверка, что боксы попадают в видимую область изображения + if x1 < 0 or y1 < 0 or x2 > width or y2 > height: + print(f"Warning: True box for class {class_names[class_id]} is out of image bounds: ({x1}, {y1}), ({x2}, {y2})") + continue + + # Вывод координат бокса в консоль для проверки + print(f"True box for class {class_names[class_id]}: ({x1}, {y1}), ({x2}, {y2})") + + # Рисуем прямоугольник и текст для истинного бокса + draw.rectangle([x1, y1, x2, y2], outline=true_color, width=3) + label = f"True: {class_names[class_id]}" + + # Добавление фона под текст для лучшей видимости + text_bbox = draw.textbbox((x1, y1 - 10), label, font=font) + draw.rectangle([text_bbox[0] - 2, text_bbox[1] - 2, text_bbox[2] + 2, text_bbox[3] + 2], fill=true_color) + draw.text((x1, y1 - 10), label, fill="white", font=font) + + # Рисуем предсказанные боксы (предполагается, что координаты уже абсолютные) + for box in pred_boxes: + class_id = int(box["cls"]) + + # Используем абсолютные координаты напрямую + x1, y1, x2, y2 = map(int, box["bbox"]) + + # Проверка, что боксы попадают в видимую область изображения + if x1 < 0 or y1 < 0 or x2 > width or y2 > height: + print(f"Warning: Predicted box for class {class_names[class_id]} is out of image bounds: ({x1}, {y1}), ({x2}, {y2})") + continue + + # Вывод координат бокса в консоль для проверки + print(f"Predicted box for class {class_names[class_id]}: ({x1}, {y1}), ({x2}, {y2})") + + # Рисуем прямоугольник и текст для предсказанного бокса + draw.rectangle([x1, y1, x2, y2], outline=pred_color, width=3) + label = f"Pred: {class_names[class_id]}" + + # Добавление фона под текст для лучшей видимости + text_bbox = draw.textbbox((x1, y1 - 10), label, font=font) + draw.rectangle([text_bbox[0] - 2, text_bbox[1] - 2, text_bbox[2] + 2, text_bbox[3] + 2], fill=pred_color) + draw.text((x1, y1 - 10), label, fill="white", font=font) + + return image + + +# Функция обработки изображений из указанной папки +def process_images_in_folder(input_folder, output_folder, conf_thres=0.4, iou_thres=0.45, use_dual_function=True): + image_folder = Path(input_folder) / "images" + label_folder = Path(input_folder) / "labels" + output_combined_folder = Path(output_folder) / "combined_boxes" + + # Создаем выходные папки, если их нет + output_combined_folder.mkdir(parents=True, exist_ok=True) + + # Выбор функции детекции + detection_function = detect_image_dual if use_dual_function else detect_image + + # Обрабатываем все изображения в папке + for image_path in image_folder.glob("*.*"): + if image_path.suffix.lower() not in [".jpg", ".jpeg", ".png"]: + continue + + image = Image.open(image_path).convert("RGB") + label_path = label_folder / (image_path.stem + ".txt") + + # Чтение истинных боксов (если необходимо) + true_boxes = [] + if label_path.exists(): + with open(label_path, 'r') as f: + for line in f: + parts = line.strip().split() + class_id = int(parts[0]) + bbox = list(map(float, parts[1:])) + true_boxes.append({"cls": class_id, "bbox": bbox}) + + # Получение предсказанных боксов + pred_boxes = detection_function( + weight_path, + image_path, + device="cpu", + conf_thres=conf_thres, + iou_thres=iou_thres, + ) + + print(pred_boxes) + + # Рисуем истинные и предсказанные боксы на одном изображении + combined_image = image.copy() + combined_image = draw_boxes_on_image(combined_image, true_boxes, pred_boxes) + combined_image.save(output_combined_folder / image_path.name) + + print("Обработка завершена.") + +# Использование функции +input_folder = r"C:\Users\user\Desktop\Проволока\снимки для КРИТБИ\yoloimages\all_data" # Путь до папки с папками images и labels +output_folder = r"C:\Users\user\Desktop\Проволока\снимки для КРИТБИ\yoloimages\out" # Путь для сохранения изображений с боксами + +process_images_in_folder(input_folder, output_folder) \ No newline at end of file diff --git a/runs_commants b/runs_commants new file mode 100644 index 000000000..7b7ae566e --- /dev/null +++ b/runs_commants @@ -0,0 +1 @@ + python train.py --batch 32 --epochs 5 --img 640 --device 0 --min-items 0 --data C:\Users\user\Desktop\data_and_weight\data\data.yaml --project C:/Users/user/Desktop/data_and_weight/ --weights C:/Users/user/Desktop/data_and_weight/gelan-c.pt --cfg models/detect/gelan-c.yaml --hyp hyp.scratch-high.yaml --workers 0 \ No newline at end of file diff --git a/split_dataset.py b/split_dataset.py new file mode 100644 index 000000000..94b7fe3b5 --- /dev/null +++ b/split_dataset.py @@ -0,0 +1,54 @@ +import os +import random +import shutil + +def split_dataset(images_dir, labels_dir, output_dir, train_ratio=0.7901, val_ratio=0.19, test_ratio=0.0201): + # Убедитесь, что пропорции суммируются до 1.0 + #assert train_ratio + val_ratio + test_ratio == 1.0, "Train, val, and test ratios must sum to 1.0" + + # Создание директорий для train, val и test + os.makedirs(os.path.join(output_dir, 'train', 'images'), exist_ok=True) + os.makedirs(os.path.join(output_dir, 'train', 'labels'), exist_ok=True) + os.makedirs(os.path.join(output_dir, 'val', 'images'), exist_ok=True) + os.makedirs(os.path.join(output_dir, 'val', 'labels'), exist_ok=True) + os.makedirs(os.path.join(output_dir, 'test', 'images'), exist_ok=True) + os.makedirs(os.path.join(output_dir, 'test', 'labels'), exist_ok=True) + + # Список всех файлов изображений + images = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))] + + # Перемешивание списка изображений + random.shuffle(images) + + # Определение количества файлов для каждой части + total_images = len(images) + train_count = int(total_images * train_ratio) + val_count = int(total_images * val_ratio) + test_count = total_images - train_count - val_count + + # Разделение файлов + train_images = images[:train_count] + val_images = images[train_count:train_count + val_count] + test_images = images[train_count + val_count:] + + # Функция для копирования файлов изображений и аннотаций + def copy_files(file_list, subset): + for image_file in file_list: + label_file = os.path.splitext(image_file)[0] + '.txt' + try: + shutil.copy(os.path.join(images_dir, image_file), os.path.join(output_dir, subset, 'images', image_file)) + shutil.copy(os.path.join(labels_dir, label_file), os.path.join(output_dir, subset, 'labels', label_file)) + except Exception as e: + print(f"Could not copy {image_file} or {label_file}: {e}") + + # Копирование файлов + copy_files(train_images, 'train') + copy_files(val_images, 'val') + copy_files(test_images, 'test') + + +# Пример использования скрипта +images_dir = r"C:\Users\user\Desktop\ddatas\data\images" +labels_dir = r"C:\Users\user\Desktop\ddatas\data\labels" +output_dir = r"C:\Users\user\Desktop\ddatas\datadatadata" +split_dataset(images_dir, labels_dir, output_dir) \ No newline at end of file diff --git a/test.py b/test.py new file mode 100644 index 000000000..28dcfc86a --- /dev/null +++ b/test.py @@ -0,0 +1,8 @@ +from detect_function import detect_image + +weights_path = r"C:\Users\user\Desktop\ddata\exp3\weights\best.pt" +image_path = r"C:\Users\user\Desktop\test_images" +results = detect_image(weights=weights_path, source=image_path) + +for res in results: + print(res) \ No newline at end of file diff --git a/test_detect.py b/test_detect.py new file mode 100644 index 000000000..4252483f8 --- /dev/null +++ b/test_detect.py @@ -0,0 +1,9 @@ +from detect_function_dual import detect_image_dual + +weight_path = r"C:\Users\pasha\OneDrive\Рабочий стол\yolo_weights\yolo_rusal.pt" +source = r"C:\Users\pasha\OneDrive\Рабочий стол\photo_2024-06-18_14-25-05.jpg" + + +res = detect_image_dual(weight_path, source, device="cpu") + +print(res) diff --git a/tg_bot.py b/tg_bot.py new file mode 100644 index 000000000..43ee095bc --- /dev/null +++ b/tg_bot.py @@ -0,0 +1,313 @@ +import os +import zipfile +import nest_asyncio +import io +import shutil +from telegram import Update, InputFile, KeyboardButton, ReplyKeyboardMarkup, ReplyKeyboardRemove +from telegram.ext import ApplicationBuilder, CommandHandler, MessageHandler, filters, ContextTypes +from PIL import Image, ImageDraw, ImageFont +from detect_function import detect_image +from telegram.error import Forbidden + +# Словарь классов и цветов +CLASS_NAMES = {0: 'Dent', 1: 'WireFlaw', 2: 'Puncture'} +CLASS_COLORS = {0: 'red', 1: 'green', 2: 'blue'} + +# Переменные для отслеживания режима и порога confidence +USER_MODES = {} +USER_CONFIDENCE = {} + +# Максимальный размер файла в байтах (1 ГБ) +MAX_FILE_SIZE = 1 * 1024 * 1024 * 1024 + +# Максимальные размеры изображения +MAX_IMAGE_SIZE = 1000, 1000 + +# Путь к весам вашей модели +MODEL_PATH = r"C:\Users\user\Desktop\crispi_defects\data\exp3\weights\best.pt" + +# Функция для обработки изображений +def process_images(image_dir, conf_thres): + results = detect_image(weights=MODEL_PATH, + source=image_dir, + conf_thres=conf_thres, + device="cpu",) + return results + +# Функция для распаковки архива +def unzip_file(file_data): + image_files = {} + with zipfile.ZipFile(io.BytesIO(file_data), 'r') as zip_ref: + for file in zip_ref.namelist(): + if file.lower().endswith(('.bmp', '.dng', '.jpeg', '.jpg', '.mpo', '.png', '.tif', '.tiff', '.webp', '.pfm')): + image_files[file] = zip_ref.read(file) + return image_files + +# Функция для изменения размера изображений +def resize_image(image): + original_size = image.size + image.thumbnail(MAX_IMAGE_SIZE, Image.LANCZOS) + return image, original_size + +# Функция для рисования боксов на изображении +def draw_boxes(image_bytes, bboxes): + image = Image.open(io.BytesIO(image_bytes)) + image, original_size = resize_image(image) + draw = ImageDraw.Draw(image) + try: + # Увеличение размера шрифта и установка кириллического шрифта + font = ImageFont.truetype("arial.ttf", 40) + except IOError: + font = ImageFont.load_default() + + scale_x = image.size[0] / original_size[0] + scale_y = image.size[1] / original_size[1] + + if not bboxes: + draw.text((10, 10), "No defects found", fill='red', font=font) + else: + for bbox in bboxes: + box = [int(coord * scale_x) if i % 2 == 0 else int(coord * scale_y) for i, coord in enumerate(bbox['bbox'])] + conf = bbox['conf'] + cls = bbox['cls'] + color = CLASS_COLORS.get(cls, 'red') + label = CLASS_NAMES.get(cls, 'Unknown') + draw.rectangle(box, outline=color, width=3) + # Добавляем контур к тексту + text = f'{label} {conf:.2f}' + text_size = draw.textbbox((0, 0), text, font=font) + text_width = text_size[2] - text_size[0] + text_height = text_size[3] - text_size[1] + x, y = box[0], box[1] - text_height + draw.rectangle([x, y, x + text_width, y + text_height], fill=color) + text_color = 'black' if color == 'yellow' else 'white' + draw.text((x, y), text, fill=text_color, font=font) + + output = io.BytesIO() + image.save(output, format='JPEG') + output.seek(0) + return output + +# Функция для нормализации координат +def normalize_bbox(bbox, width, height): + x_min, y_min, x_max, y_max = bbox + x_center = (x_min + x_max) / 2.0 / width + y_center = (y_min + y_max) / 2.0 / height + box_width = (x_max - x_min) / width + box_height = (y_max - y_min) / height + return x_center, y_center, box_width, box_height + +# Функция для создания архива с результатами +def create_results_archive(results, image_files, original_filename): + archive_buffer = io.BytesIO() + with zipfile.ZipFile(archive_buffer, 'w', zipfile.ZIP_DEFLATED) as archive: + for image_path, bboxes in results.items(): + img_bytes = image_files[image_path] + with Image.open(io.BytesIO(img_bytes)) as img: + img = resize_image(img)[0] + width, height = img.size + txt_content = "" + for bbox in bboxes: + x_center, y_center, box_width, box_height = normalize_bbox(bbox['bbox'], width, height) + conf = bbox['conf'] + cls = bbox['cls'] + txt_content += f'{cls} {x_center:.6f} {y_center:.6f} {box_width:.6f} {box_height:.6f} {conf:.6f}\n' + txt_filename = os.path.splitext(image_path)[0] + '.txt' + archive.writestr(txt_filename, txt_content) + archive_buffer.seek(0) + archive_buffer.name = f"{os.path.splitext(original_filename)[0]}_labels.zip" + return archive_buffer + +# Функция для создания CSV с результатами +def create_results_csv(results, image_files): + csv_buffer = io.StringIO() + csv_buffer.write("filename;class_id;rel_x;rel_y;width;height\n") + for image_path, bboxes in results.items(): + img_bytes = image_files[image_path] + with Image.open(io.BytesIO(img_bytes)) as img: + img = resize_image(img)[0] + width, height = img.size + for bbox in bboxes: + x_center, y_center, box_width, box_height = normalize_bbox(bbox['bbox'], width, height) + cls = bbox['cls'] + csv_buffer.write(f"{os.path.basename(image_path)};{cls};{x_center:.6f};{y_center:.6f};{box_width:.6f};{box_height:.6f}\n") + csv_buffer.seek(0) + return io.BytesIO(csv_buffer.getvalue().encode('utf-8')) + +# Функция для создания клавиатуры +def get_keyboard(): + keyboard = [ + [KeyboardButton("Выбрать пороговое значение")] + ] + return ReplyKeyboardMarkup(keyboard, one_time_keyboard=False, resize_keyboard=True) + +# Обработчик команды /start +async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + USER_MODES[update.effective_user.id] = None + USER_CONFIDENCE[update.effective_user.id] = 0.1 # Значение по умолчанию + await update.message.reply_text('Привет! Отправьте мне фото, изображение или архив с изображениями для обработки.', reply_markup=get_keyboard()) + +# Функция для подсчета дефектов по классам +def count_defects(bboxes): + counts = {CLASS_NAMES[cls]: 0 for cls in CLASS_NAMES} + for bbox in bboxes: + cls = bbox['cls'] + counts[CLASS_NAMES[cls]] += 1 + total = sum(counts.values()) + return counts, total + +# Обработчик полученных файлов и сообщений +async def handle_file(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + user_id = update.message.from_user.id + conf_thres = USER_CONFIDENCE.get(user_id, 0.1) + + try: + if update.message.photo: + await update.message.reply_text('Изображение получено, подождите немного.') + file = await context.bot.get_file(update.message.photo[-1].file_id) + photo_bytes = await file.download_as_bytearray() + + # Обрабатываем фото + image_dir = 'temp_image' + os.makedirs(image_dir, exist_ok=True) + with open(os.path.join(image_dir, 'photo.jpg'), 'wb') as img_file: + img_file.write(photo_bytes) + + detection_results = process_images(image_dir, conf_thres) + results = {'photo.jpg': detection_results} + counts, total = count_defects(detection_results) + + # Отправляем фото с нарисованными боксами и статистику + img_bytes = draw_boxes(photo_bytes, detection_results) + if total == 0: + await context.bot.send_photo(chat_id=update.message.chat_id, photo=img_bytes, caption="Дефектов не найдено.") + else: + caption = "\n".join([f"{cls}: {count}" for cls, count in counts.items()]) + f"\nВсего: {total}" + await context.bot.send_photo(chat_id=update.message.chat_id, photo=img_bytes, caption=caption) + + # Удаляем временные файлы + shutil.rmtree(image_dir) + + elif update.message.document: + if update.message.document.file_size > MAX_FILE_SIZE: + await update.message.reply_text('Файл слишком большой. Максимальный размер файла - 1 ГБ.') + return + + if update.message.document.mime_type.startswith('image/'): + await update.message.reply_text('Изображение получено, подождите немного.') + file = await context.bot.get_file(update.message.document.file_id) + photo_bytes = await file.download_as_bytearray() + + # Обрабатываем фото + image_dir = 'temp_image' + os.makedirs(image_dir, exist_ok=True) + with open(os.path.join(image_dir, 'photo.jpg'), 'wb') as img_file: + img_file.write(photo_bytes) + + detection_results = process_images(image_dir, conf_thres) + results = {'photo.jpg': detection_results} + counts, total = count_defects(detection_results) + + # Отправляем фото с нарисованными боксами и статистику + img_bytes = draw_boxes(photo_bytes, detection_results) + if total == 0: + await context.bot.send_photo(chat_id=update.message.chat_id, photo=img_bytes, caption="Дефектов не найдено.") + else: + caption = "\n".join([f"{cls}: {count}" for cls, count in counts.items()]) + f"\nВсего: {total}" + await context.bot.send_photo(chat_id=update.message.chat_id, photo=img_bytes, caption=caption) + + # Удаляем временные файлы + shutil.rmtree(image_dir) + + elif update.message.document.mime_type == 'application/zip': + await update.message.reply_text('Архив получен, подождите немного.') + file = await context.bot.get_file(update.message.document.file_id) + file_data = await file.download_as_bytearray() + original_filename = update.message.document.file_name + + # Распаковываем архив + image_files = unzip_file(file_data) + + if not image_files: + await update.message.reply_text('Не найдено изображений поддерживаемых форматов в архиве.') + return + + # Сохраняем изображения в временную директорию для обработки + image_dir = 'images' + if os.path.exists(image_dir): + shutil.rmtree(image_dir) + os.makedirs(image_dir) + for img_name, img_data in image_files.items(): + with open(os.path.join(image_dir, os.path.basename(img_name)), 'wb') as img_file: + img_file.write(img_data) + + # Обрабатываем изображения + detection_results = process_images(image_dir, conf_thres) + + # Преобразуем результаты в нужный формат + results = {} + for result in detection_results: + image_path = os.path.basename(result['path']) + if image_path not in results: + results[image_path] = [] + results[image_path].append({ + 'bbox': result['bbox'], + 'conf': result['conf'], + 'cls': result['cls'] + }) + + # Формируем и отправляем CSV с результатами + results_csv = create_results_csv(results, image_files) + await context.bot.send_document(chat_id=update.message.chat_id, document=InputFile(results_csv, filename="submission.csv")) + + # Удаляем временные файлы + shutil.rmtree(image_dir) + + else: + await update.message.reply_text('Пожалуйста, отправьте фото, изображение или архив с изображениями.') + + else: + await update.message.reply_text('Пожалуйста, отправьте фото, изображение или архив с изображениями.') + + except Forbidden: + print(f"Bot was blocked by the user: {update.message.chat_id}") + except Exception as e: + await update.message.reply_text(f'Произошла ошибка: {e}') + +# Обработчик текстовых сообщений +async def handle_text(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + user_id = update.message.from_user.id + text = update.message.text + if text == "Выбрать пороговое значение": + await update.message.reply_text('Введите значение порога от 1 до 99:', reply_markup=ReplyKeyboardRemove()) + USER_MODES[user_id] = 'set_threshold' + elif USER_MODES.get(user_id) == 'set_threshold': + try: + value = int(text) + if 1 <= value <= 99: + USER_CONFIDENCE[user_id] = value / 100.0 + await update.message.reply_text(f'Установлен порог confidence: {USER_CONFIDENCE[user_id]}', reply_markup=get_keyboard()) + USER_MODES[user_id] = None + else: + await update.message.reply_text('Пожалуйста, введите значение от 1 до 99.') + except ValueError: + await update.message.reply_text('Пожалуйста, введите корректное числовое значение.') + else: + await update.message.reply_text('Пожалуйста, отправьте фото, изображение или архив с изображениями.') + +def main() -> None: + nest_asyncio.apply() # Обходим проблему с уже запущенным event loop + + # Вставьте сюда ваш токен от BotFather + token = 'YOUR_BOT_TOKEN_HERE' + + application = ApplicationBuilder().token(token).build() + + application.add_handler(CommandHandler("start", start)) + application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_text)) + application.add_handler(MessageHandler(filters.Document.ALL | filters.PHOTO, handle_file)) + + application.run_polling() + +if __name__ == '__main__': + main() diff --git a/train_notebook.ipynb b/train_notebook.ipynb new file mode 100644 index 000000000..060f664b2 --- /dev/null +++ b/train_notebook.ipynb @@ -0,0 +1,77 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\u001b[34m\u001b[1mtrain: \u001b[0mweights={HOME}/weights/gelan-c.pt, cfg=models/detect/gelan-c.yaml, data=C:/Users/user/Desktop/fissures_bubbles_data/fissures_bubbles_data/data.yaml, hyp=hyp.scratch-high.yaml, epochs=300, batch_size=16, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, noplots=False, evolve=None, bucket=, cache=None, image_weights=False, device=0, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=runs\\train, name=exp, exist_ok=False, quad=False, cos_lr=False, flat_cos_lr=False, fixed_lr=False, label_smoothing=0.0, patience=100, freeze=[0], save_period=-1, seed=0, local_rank=-1, min_items=0, close_mosaic=50, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest\n", + "YOLO 519346b Python-3.10.7 torch-1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3090 Ti, 24564MiB)\n", + "\n", + "\u001b[34m\u001b[1mhyperparameters: \u001b[0mlr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, cls_pw=1.0, obj=0.7, obj_pw=1.0, dfl=1.5, iou_t=0.2, anchor_t=5.0, fl_gamma=0.0, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.9, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.15, copy_paste=0.3\n", + "\u001b[34m\u001b[1mClearML: \u001b[0mrun 'pip install clearml' to automatically track, visualize and remotely train YOLO in ClearML\n", + "\u001b[34m\u001b[1mComet: \u001b[0mrun 'pip install comet_ml' to automatically track and visualize YOLO runs in Comet\n", + "\u001b[34m\u001b[1mTensorBoard: \u001b[0mStart with 'tensorboard --logdir runs\\train', view at http://localhost:6006/\n", + "\n", + "Dataset not found , missing paths ['C:\\\\content\\\\drive\\\\My Drive\\\\fissures_bubbles_data\\\\valid\\\\images']\n", + "Traceback (most recent call last):\n", + " File \"c:\\Users\\user\\yolo9_train\\yolov9\\train.py\", line 634, in \n", + " main(opt)\n", + " File \"c:\\Users\\user\\yolo9_train\\yolov9\\train.py\", line 528, in main\n", + " train(opt.hyp, opt, device, callbacks)\n", + " File \"c:\\Users\\user\\yolo9_train\\yolov9\\train.py\", line 95, in train\n", + " data_dict = data_dict or check_dataset(data) # check if None\n", + " File \"c:\\Users\\user\\yolo9_train\\yolov9\\utils\\general.py\", line 537, in check_dataset\n", + " raise Exception('Dataset not found ❌')\n", + "Exception: Dataset not found ❌\n" + ] + } + ], + "source": [ + "!python train.py \n", + "--batch 16 --epochs 300 --img 640 --device 0 --min-items 0 --close-mosaic 50 \n", + "--data C:/Users/user/Desktop/fissures_bubbles_data/fissures_bubbles_data/data.yaml \n", + "--weights C:/Users/user/Desktop/gelan-c.pt \n", + "--cfg models/detect/gelan-c.yaml \n", + "--hyp hyp.scratch-high.yaml \n", + "\n", + "\n", + "\n", + "\n", + "docker run --gpus all -it -v C:/Users/user/Desktop/fissures_bubbles_data/fissures_bubbles_data yolo9\n", + "\n", + "python train.py \n", + "--batch 16 --epochs 300 --img 640 --device 0 --min-items 0\n", + "--data path/to//data.yaml \n", + "--weights path/to/weights.pt\n", + "--cfg models/detect/gelan-c.yaml \n", + "--hyp hyp.scratch-high.yaml " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.7" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/utils/dataloaders.py b/utils/dataloaders.py index 776042999..4f95471d7 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -334,6 +334,36 @@ def _cv2_rotate(self, im): def __len__(self): return self.nf # number of files +class LoadCV2Image: + def __init__(self, image, img_size=640, stride=32, auto=True, transforms=None): + self.img_size = img_size + self.stride = stride + self.files = [image] + self.nf = 1 # number of files + self.mode = 'image' + self.auto = auto + self.transforms = transforms # optional + + def __iter__(self): + self.count = 0 + return self + + def __next__(self): + if self.count == self.nf: + raise StopIteration + im0 = self.files[self.count] # BGR + s = f'image {self.count + 1}/{self.nf}' + + if self.transforms: + im = self.transforms(im0) # transforms + else: + im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0] # padded resize + im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB + im = np.ascontiguousarray(im) # contiguous + + self.count += 1 + return s, im, im0, None, s + class LoadStreams: # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams`