add files for web demo

app.py

@@ -0,0 +1,318 @@
+# import gradio as gr
+# import torch
+# import torch.nn as nn
+# from torchvision import models, transforms
+# from PIL import Image
+# from transformers import AutoModel, AutoTokenizer
+# import easyocr
+# import json
+# import os
+
+# import spaces  # добавьте в начале
+
+# @spaces.GPU(duration=60)  # добавьте перед predict
+# def predict_demo(image, caption_text=""):
+#     # ... ваш код
+# # ======================
+# # ФИКСИРУЕМ ПУТИ (важно для Spaces!)
+# # ======================
+# # Модели и веса лежат в той же папке, что и app.py
+# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# # Загрузка названий классов
+# with open(os.path.join(BASE_DIR, "class_names.json"), "r") as f:
+#     id2label = json.load(f)
+#     id2label = {int(k): v for k, v in id2label.items()}
+
+# NUM_CLASSES = len(id2label)
+
+
+# # ======================
+# # ЗАГРУЗКА МОДЕЛЕЙ (один раз, с кешированием)
+# # ======================
+# @gr.cache_resource
+# def load_models():
+#     DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#     print(f"Using device: {DEVICE}")
+
+#     # Визуальный энкодер
+#     visual = models.resnet50(weights=None)
+#     visual.fc = nn.Identity()
+#     visual.load_state_dict(torch.load(os.path.join(BASE_DIR, "resnet50_encoder.pth"), map_location=DEVICE))
+#     visual.to(DEVICE)
+#     visual.eval()
+#     for p in visual.parameters():
+#         p.requires_grad = False
+
+#     # Текстовые энкодеры
+#     tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
+#     ocr_encoder = AutoModel.from_pretrained("cointegrated/rubert-tiny2").to(DEVICE).eval()
+#     caption_encoder = AutoModel.from_pretrained("cointegrated/rubert-tiny2").to(DEVICE).eval()
+
+#     for p in ocr_encoder.parameters():
+#         p.requires_grad = False
+#     for p in caption_encoder.parameters():
+#         p.requires_grad = False
+
+#     # Классификатор
+#     class ConcatFusionModel(nn.Module):
+#         def __init__(self, num_classes, dropout=0.3):
+#             super().__init__()
+#             self.classifier = nn.Sequential(
+#                 nn.Linear(2048 + 312 + 312, 512),
+#                 nn.BatchNorm1d(512),
+#                 nn.ReLU(),
+#                 nn.Dropout(dropout),
+#                 nn.Linear(512, num_classes)
+#             )
+
+#         def forward(self, v, ocr, cap):
+#             x = torch.cat([v, ocr, cap], dim=1)
+#             return self.classifier(x)
+
+#     model = ConcatFusionModel(NUM_CLASSES, dropout=0.3)
+#     model.load_state_dict(torch.load(os.path.join(BASE_DIR, "best_concat_model.pth"), map_location=DEVICE))
+#     model.to(DEVICE)
+#     model.eval()
+
+#     # EasyOCR
+#     reader = easyocr.Reader(["ru", "en"], gpu=(DEVICE.type == "cuda"))
+
+#     # Трансформы
+#     val_transform = transforms.Compose([
+#         transforms.Resize(256),
+#         transforms.CenterCrop(224),
+#         transforms.ToTensor(),
+#         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+#     ])
+
+#     return visual, ocr_encoder, caption_encoder, tokenizer, model, reader, val_transform, DEVICE
+
+
+# # Загружаем всё при старте
+# visual, ocr_encoder, caption_encoder, tokenizer, model, reader, val_transform, DEVICE = load_models()
+
+
+# # ======================
+# # ФУНКЦИЯ ПРЕДСКАЗАНИЯ
+# # ======================
+# def predict(image, caption_text=""):
+#     image = image.convert("RGB")
+
+#     # OCR
+#     ocr_result = reader.readtext(np.array(image), detail=0, paragraph=True)
+#     ocr_text = " ".join(ocr_result) if ocr_result else ""
+
+#     # Image
+#     image_tensor = val_transform(image).unsqueeze(0).to(DEVICE)
+#     with torch.no_grad():
+#         v = visual(image_tensor)
+#         v = torch.flatten(v, 1)
+
+#     # OCR encode
+#     ocr_enc = tokenizer(ocr_text, truncation=True, padding="max_length", max_length=64, return_tensors="pt")
+#     ocr_ids = ocr_enc["input_ids"].to(DEVICE)
+#     ocr_mask = ocr_enc["attention_mask"].to(DEVICE)
+#     with torch.no_grad():
+#         ocr_out = ocr_encoder(input_ids=ocr_ids, attention_mask=ocr_mask)
+#         ocr = ocr_out.last_hidden_state[:, 0]
+
+#     # Caption encode
+#     cap_enc = tokenizer(caption_text, truncation=True, padding="max_length", max_length=128, return_tensors="pt")
+#     cap_ids = cap_enc["input_ids"].to(DEVICE)
+#     cap_mask = cap_enc["attention_mask"].to(DEVICE)
+#     with torch.no_grad():
+#         cap_out = caption_encoder(input_ids=cap_ids, attention_mask=cap_mask)
+#         cap = cap_out.last_hidden_state[:, 0]
+
+#     # Предсказание
+#     with torch.no_grad():
+#         logits = model(v, ocr, cap)
+#         probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
+
+#     result = {id2label[i]: float(probs[i]) for i in range(NUM_CLASSES)}
+#     return dict(sorted(result.items(), key=lambda x: x[1], reverse=True))
+
+
+# # ======================
+# # GRADIO ИНТЕРФЕЙС
+# # ======================
+# demo = gr.Interface(
+#     fn=predict,
+#     inputs=[
+#         gr.Image(type="pil", label="Загрузите изображение"),
+#         gr.Textbox(label="Подпись (необязательно)", placeholder="Введите текст подписи...")
+#     ],
+#     outputs=gr.Label(num_top_classes=5, label="Предсказанные категории"),
+#     title="Мультимодальный классификатор контента",
+#     description="Модель анализирует изображение + подпись + текст на картинке (EasyOCR)"
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+
+
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+import easyocr
+import json
+import os
+import numpy as np
+
+import spaces
+
+# ======================
+# УСТАНОВКА УСТРОЙСТВА
+# ======================
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {DEVICE}")
+
+# ======================
+# ПУТИ
+# ======================
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Загрузка названий классов
+with open(os.path.join(BASE_DIR, "class_names.json"), "r") as f:
+    id2label = json.load(f)
+    id2label = {int(k): v for k, v in id2label.items()}
+
+NUM_CLASSES = len(id2label)
+
+
+# ======================
+# ОПРЕДЕЛЕНИЕ МОДЕЛИ
+# ======================
+class ConcatFusionModel(nn.Module):
+    def __init__(self, num_classes, dropout=0.3):
+        super().__init__()
+        self.classifier = nn.Sequential(
+            nn.Linear(2048 + 312 + 312, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, num_classes)
+        )
+
+    def forward(self, v, ocr, cap):
+        x = torch.cat([v, ocr, cap], dim=1)
+        return self.classifier(x)
+
+
+# ======================
+# ЗАГРУЗКА МОДЕЛЕЙ
+# ======================
+@gr.cache_resource
+def load_models():
+    # Визуальный энкодер (загружаем предобученный из torchvision)
+    visual = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+    visual.fc = nn.Identity()  # убираем классификатор
+    visual.to(DEVICE)
+    visual.eval()
+    for p in visual.parameters():
+        p.requires_grad = False
+
+    # Текстовые энкодеры (загружаем предобученные из Hugging Face)
+    tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
+    ocr_encoder = AutoModel.from_pretrained("cointegrated/rubert-tiny2").to(DEVICE).eval()
+    caption_encoder = AutoModel.from_pretrained("cointegrated/rubert-tiny2").to(DEVICE).eval()
+
+    for p in ocr_encoder.parameters():
+        p.requires_grad = False
+    for p in caption_encoder.parameters():
+        p.requires_grad = False
+
+    # Классификационная голова (обученная)
+    model = ConcatFusionModel(NUM_CLASSES, dropout=0.3)
+    model.load_state_dict(torch.load(os.path.join(BASE_DIR, "concat_model.pth"), map_location=DEVICE))
+    model.to(DEVICE)
+    model.eval()
+
+    # EasyOCR
+    reader = easyocr.Reader(["ru", "en"], gpu=(DEVICE.type == "cuda"))
+
+    # Трансформы для изображений
+    val_transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    return visual, ocr_encoder, caption_encoder, tokenizer, model, reader, val_transform
+
+
+visual, ocr_encoder, caption_encoder, tokenizer, model, reader, val_transform = load_models()
+
+
+# ======================
+# ФУНКЦИЯ ПРЕДСКАЗАНИЯ
+# ======================
+@spaces.GPU(duration=60)
+def predict(image, caption_text=""):
+    image = image.convert("RGB")
+
+    # OCR
+    ocr_result = reader.readtext(np.array(image), detail=0, paragraph=True)
+    ocr_text = " ".join(ocr_result) if ocr_result else ""
+
+    # Image
+    image_tensor = val_transform(image).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        v = visual(image_tensor)
+        v = torch.flatten(v, 1)
+
+    # OCR encode
+    ocr_enc = tokenizer(ocr_text, truncation=True, padding="max_length", max_length=64, return_tensors="pt")
+    with torch.no_grad():
+        ocr_out = ocr_encoder(
+            input_ids=ocr_enc["input_ids"].to(DEVICE),
+            attention_mask=ocr_enc["attention_mask"].to(DEVICE)
+        )
+        ocr = ocr_out.last_hidden_state[:, 0]
+
+    # Caption encode
+    cap_enc = tokenizer(caption_text, truncation=True, padding="max_length", max_length=128, return_tensors="pt")
+    with torch.no_grad():
+        cap_out = caption_encoder(
+            input_ids=cap_enc["input_ids"].to(DEVICE),
+            attention_mask=cap_enc["attention_mask"].to(DEVICE)
+        )
+        cap = cap_out.last_hidden_state[:, 0]
+
+    # Предсказание
+    with torch.no_grad():
+        logits = model(v, ocr, cap)
+        probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
+
+    result = {id2label[i]: float(probs[i]) for i in range(NUM_CLASSES)}
+    return dict(sorted(result.items(), key=lambda x: x[1], reverse=True))
+
+
+# ======================
+# GRADIO ИНТЕРФЕЙС
+# ======================
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="📸 Загрузите изображение"),
+        gr.Textbox(label="📝 Подпись (необязательно)", placeholder="Введите текст подписи...")
+    ],
+    outputs=gr.Label(num_top_classes=5, label="🎯 Предсказанные категории"),
+    title="Мультимодальный классификатор контента",
+    description="Модель анализирует изображение + подпись + текст на картинке (EasyOCR)"
+)
+
+if __name__ == "__main__":
+    demo.launch()

class_names.json

@@ -0,0 +1,10 @@
+{
+  "0": "Животные",
+  "1": "Кулинария",
+  "2": "Путешествия",
+  "3": "Развлечения и юмор",
+  "4": "СМИ",
+  "5": "Торговля и объявления",
+  "6": "Увлечения и хобби",
+  "7": "Философия и религия"
+}

concat_model.pth

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

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+torch
+torchvision
+transformers
+easyocr
+Pillow
+numpy
+scikit-learn