init

CoinCounter/model.py

@@ -0,0 +1,24 @@
+import torch
+from ultralytics import YOLO
+from .utils import Result
+
+# Загрузка модели YOLO
+model = YOLO('CoinCounter/models/model.onnx', task='detect',verbose=False)
+
+def predict(path='data', conf=0.5, iou=0.5):
+    """
+    Выполняет предсказание с помощью модели YOLO.
+
+    :param path: Путь к изображению или директории с изображениями для предсказания.
+    :param conf: Порог уверенности для предсказаний (0.0 - 1.0). Чем выше значение, тем меньше ложных срабатываний.
+    :param iou: Порог перекрытия для фильтрации предсказанных рамок (0.0 - 1.0). Чем выше значение, тем более строгий отбор.
+    :return: Результаты предсказания в формате объекта Result.
+    """
+    # Выполнение предсказания с заданными параметрами
+    results = model.predict(
+        path, 
+        conf=conf,   # Параметр уверенности
+        iou=iou,     # Параметр IoU
+        verbose=False
+    )
+    return Result(results)

CoinCounter/models/model.onnx

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

CoinCounter/utils.py

@@ -0,0 +1,132 @@
+import pandas as pd
+import matplotlib.pyplot as plt
+from PIL import Image, ImageDraw
+import numpy as np
+import matplotlib.colors as mcolors
+
+class_map = {
+    0: 'Dime',
+    1: 'Nickel',
+    2: 'Penny',
+    3: 'Quarter',
+    4: 'fifty',
+    5: 'five',
+    6: 'hundred',
+    7: 'one',
+    8: 'ten',
+    9: 'twenty'
+}
+
+value_map = {
+    'Penny': 0.01,
+    'Nickel': 0.05,
+    'Dime': 0.10,
+    'Quarter': 0.25,
+    'one': 1.00,
+    'five': 5.00,
+    'ten': 10.00,
+    'twenty': 20.00,
+    'fifty': 50.00,
+    'hundred': 100.00
+}
+
+nomimals = ['Dime', 'Nickel', 'Penny', 'Quarter', 'fifty', 'five', 'hundred', 'one', 'ten', 'twenty']
+
+class Result:
+    def __init__(self, results, nomimals = nomimals, class_map = class_map, value_map = value_map):
+        self.results = results
+        self.nomimals = nomimals
+        self.class_map = class_map
+        self.value_map = value_map
+
+        # Список объектов по классам
+        self.objects = [[int(box.cls) for box in result.boxes] for result in self.results]
+        
+        # Создаем DataFrame с подсчетом объектов и их общей "стоимостью"
+        self.df = pd.DataFrame([
+            {class_map[obj]: lst.count(obj) for obj in class_map} for lst in self.objects
+        ])
+        self.df['total'] = self.df.apply(lambda row: sum(row[obj] * value_map[obj] for obj in value_map), axis=1)
+
+    def __len__(self):
+        return len(self.objects)
+
+    def total(self, id=None):
+        if id is None:
+            return round(self.df.total.sum(), 2)
+        elif isinstance(id, int):
+            return round(self.df.iloc[id].total, 2)
+        else:
+            print('Please use an integer for the ID.')
+
+    def image(self, index=0):
+        """
+        Возвращает изображение с разметкой объектов.
+        :param index: Индекс изображения в results.
+        :return: PIL.Image с размеченными объектами.
+        """
+        if index >= len(self.results):
+            raise IndexError("Индекс выходит за пределы диапазона.")
+
+        # Получаем изображение и разметку
+        result = self.results[index]
+        img = result.orig_img
+        boxes = result.boxes
+
+        # Конвертируем изображение из BGR в RGB
+        img_rgb = img[..., ::-1]  # Меняем порядок каналов с BGR на RGB
+
+        # Конвертируем изображение в PIL.Image
+        img_pil = Image.fromarray(img_rgb)
+        draw = ImageDraw.Draw(img_pil)
+
+        # Определяем цвета для объектов
+        object_colors = list(mcolors.TABLEAU_COLORS.values())
+
+        # Если есть разметка, добавляем её
+        if boxes is not None:
+            xyxy = boxes.xyxy.cpu().numpy()  # Координаты [x1, y1, x2, y2]
+            conf = boxes.conf.cpu().numpy()   # Уверенность
+            cls = boxes.cls.cpu().numpy()     # Классы
+
+            for j in range(len(xyxy)):
+                x1, y1, x2, y2 = xyxy[j]
+                class_id = int(cls[j])
+                class_name = result.names[class_id]
+                confidence = conf[j]
+                
+                # Присваиваем цвет для класса
+                color = object_colors[class_id % len(object_colors)]
+                
+                # Рисуем прямоугольник
+                draw.rectangle([(x1, y1), (x2, y2)], outline=color, width=3)
+                
+                # Добавляем подпись (класс и уверенность)
+                text = f'{class_name} {confidence:.2f}'
+                draw.text((x1, y1 - 10), text, fill=color)
+
+        return img_pil
+    def show(self, num_cols=1, scale=1):
+        """
+        Функция для создания коллажа с размеченными объектами на фотографиях.
+        :param num_cols: Количество столбцов в коллаже
+        :param scale: Масштаб для отображения изображений
+        """
+        num_images = len(self.results)
+        num_rows = (num_images + num_cols - 1) // num_cols  # Округление вверх
+        
+        fig, axes = plt.subplots(num_rows, num_cols, figsize=(num_cols * 5 * scale, num_rows * 5 * scale))
+        axes = axes.flatten() if isinstance(axes, (list, np.ndarray)) else [axes]
+
+        for i, ax in enumerate(axes[:num_images]):
+            img_pil = self.image(i)  # Получаем изображение с разметкой
+            ax.imshow(img_pil)
+            ax.axis('off')
+
+        # Убираем лишние оси
+        for ax in axes[num_images:]:
+            ax.axis('off')
+
+        plt.tight_layout()
+        plt.show()
+

app.py

@@ -0,0 +1,64 @@
+import os
+import gradio as gr
+from CoinCounter.model import predict
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Получаем все пути к изображениям в папке 'data'
+image_paths = [os.path.join('data', filename) for filename in os.listdir('data') if filename.endswith(('.png', '.jpg', '.jpeg'))]
+
+def process_image(image, conf, iou):
+    # Сохраняем изображение во временный буфер для обработки
+    image_path = "temp_image.jpg"
+    image.save(image_path)
+    
+    # Выполняем предсказание с моделью с заданными параметрами
+    results = predict(path=image_path, conf=conf, iou=iou)
+    
+    # Получаем изображение с размеченными объектами
+    annotated_image = results.image(0)
+    
+    # Конвертируем результат в формат, подходящий для вывода
+    annotated_image = np.array(annotated_image)
+    
+    # Получаем сумму денег
+    total_sum = results.total(0)
+    
+    # Извлекаем количество монет из DataFrame
+    df = results.df
+    coin_counts = df.iloc[0, :-1]  # исключаем последнюю колонку 'total'
+    
+    # Формируем строку с количеством монет
+    coin_info = "\n".join([f"{coin}: {int(count)}" for coin, count in zip(coin_counts.index, coin_counts) if count > 0])
+    
+    # Формируем итоговый текст
+    result_text = f"Общая сумма на изображении: {total_sum:.2f} $\n\nКоличество объектов:\n{coin_info}"
+    
+    return annotated_image, result_text
+
+# Определяем интерфейс Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("## Coin Counter")
+    
+    with gr.Row():
+        # Левый блок (ввод)
+        with gr.Column():
+            input_image = gr.Image(label="Загрузите изображение", type="pil")
+            conf_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.3, label="Порог уверенности (conf)")
+            iou_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.45, label="Порог перекрытия (IoU)")
+        
+        # Правый блок (вывод)
+        with gr.Column():
+            output_image = gr.Image(label="Размеченное изображение")
+            output_text = gr.Textbox(label="Результат", lines=5)
+    
+    # Обработка изображения автоматически при изменении изображения или тумблеров
+    input_image.change(process_image, inputs=[input_image, conf_slider, iou_slider], outputs=[output_image, output_text])
+    conf_slider.change(process_image, inputs=[input_image, conf_slider, iou_slider], outputs=[output_image, output_text])
+    iou_slider.change(process_image, inputs=[input_image, conf_slider, iou_slider], outputs=[output_image, output_text])
+    
+    # Добавление примеров из папки 'data'
+    gr.Examples(examples=image_paths, inputs=input_image)
+
+# Запускаем интерфейс
+demo.launch()

main.py

@@ -0,0 +1,71 @@
+import tkinter as tk
+from tkinter import filedialog, messagebox
+from CoinCounter.model import predict
+import os
+import glob
+
+class CoinCounterApp:
+    def __init__(self, root):
+        self.root = root
+        self.root.title("CoinCounter")
+        self.root.geometry("300x150") 
+
+        # Кнопка для выбора файла
+        self.file_button = tk.Button(root, text="Выбрать файл", command=self.select_file)
+        self.file_button.pack(pady=10)
+
+        # Кнопка для выбора папки
+        self.folder_button = tk.Button(root, text="Выбрать папку", command=self.select_folder)
+        self.folder_button.pack(pady=10)
+
+    def select_file(self):
+        file_path = filedialog.askopenfilename(
+            title="Выберите файл",
+            filetypes=[("Image files", "*.jpg *.jpeg *.png")]
+        )
+        if file_path:
+            self.process_image(file_path)
+
+    def select_folder(self):
+        folder_path = filedialog.askdirectory(title="Выберите папку")
+        if folder_path:
+            if self.check_images_in_folder(folder_path):
+                self.process_image(folder_path)
+            else:
+                messagebox.showinfo("Информация", "В папке нет изображений")
+
+    def check_images_in_folder(self, folder_path):
+        image_files = glob.glob(os.path.join(folder_path, "*.jpg")) + \
+                      glob.glob(os.path.join(folder_path, "*.jpeg")) + \
+                      glob.glob(os.path.join(folder_path, "*.png"))
+        return len(image_files) > 0
+
+    def process_image(self, path):
+        # Показать окно с сообщением о процессе прогноза
+        self.processing_window = tk.Toplevel(self.root)
+        self.processing_window.title("Процесс")
+        self.processing_window.geometry("350x100")  # Увеличиваем размер окна
+        tk.Label(self.processing_window, text="Идет процесс прогноза... пожалуйста, подождите", wraplength=300).pack(pady=30)
+        self.root.update_idletasks()  # Обновить интерфейс перед выполнением длительной задачи
+
+        self.root.after(100, self.run_prediction, path)  # Запуск функции предсказания через 100 мс
+
+    def run_prediction(self, path):
+        try:
+            # Запуск функции predict
+            results = predict(
+                path=path,
+                conf=0.3,
+                iou=0.3
+            )
+            total_amount = results.total()
+            self.processing_window.destroy()  # Закрыть окно с процессом
+            messagebox.showinfo("Результат", f"Общая сумма: {total_amount:.2f}$")
+        except Exception as e:
+            self.processing_window.destroy()  # Закрыть окно с процессом
+            messagebox.showerror("Ошибка", str(e))
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    app = CoinCounterApp(root)
+    root.mainloop()