Update app.py

app.py

@@ -7,11 +7,18 @@ import numpy as np
 from pydub import AudioSegment
 from langdetect import detect
 from rembg import remove
-import io
+import torch
+import torchvision
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
+from torchvision.transforms import functional as F
 
 # Инициализация клиента для модели HuggingFaceH4/zephyr-7b-beta
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+# Загрузка предобученной модели для обнаружения объектов
+model = fasterrcnn_resnet50_fpn(pretrained=True)
+model.eval()
+
 # Функция для обработки PDF
 def process_pdf(file):
     pdf_reader = PyPDF2.PdfReader(file)
@@ -51,6 +58,23 @@ def remove_background(image):
     output = remove(image)
     return output
 
+# Функция для определения количества объектов на изображении
+def count_objects(image):
+    if image is None:
+        return "Изображение не загружено."
+    
+    # Преобразование изображения в формат, подходящий для модели
+    img = Image.open(image.name).convert("RGB")
+    img_tensor = F.to_tensor(img).unsqueeze(0)
+    
+    # Обнаружение объектов
+    with torch.no_grad():
+        predictions = model(img_tensor)
+    
+    # Подсчет количества объектов
+    num_objects = len(predictions[0]['labels'])
+    return f"Количество объектов на изображении: {num_objects}"
+
 # Функция для определения языка текста
 def detect_language(text):
     try:
@@ -158,6 +182,14 @@ with gr.Blocks() as demo:
             remove_bg_button = gr.Button("Удалить фон")
             remove_bg_button.click(fn=remove_background, inputs=image_input, outputs=image_output)
 
+        # Вкладка Numage
+        with gr.Tab("Numage"):
+            gr.Markdown("## Numage: Подсчет объектов на изображении")
+            numage_input = gr.File(label="Загрузите изображение", file_types=["image"])
+            numage_output = gr.Textbox(label="Результат", interactive=False)
+            numage_button = gr.Button("Определить количество объектов")
+            numage_button.click(fn=count_objects, inputs=numage_input, outputs=numage_output)
+
     # Интерфейс чата
     chat_interface = gr.ChatInterface(
         respond,