Create app.py

app.py

@@ -0,0 +1,175 @@
+import os
+import io
+import torch
+import pandas as pd
+import gradio as gr
+from PIL import Image
+from sd_parsers import ParserManager
+from torchvision import transforms
+from transformers import CLIPProcessor, CLIPModel, Blip2Processor, Blip2ForConditionalGeneration
+import lpips
+import piq
+from torchmetrics import ClippedCosineSimilarity
+import plotly.express as px
+import plotly.graph_objects as go
+
+# --------------------
+# Setup Models
+# --------------------
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# CLIP for prompt alignment & aesthetics
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+# BLIP-2 for caption generation
+blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl").to(device)
+blip_model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-flan-t5-xl", torch_dtype=torch.float16
+).to(device)
+
+# LPIPS for diversity
+lpips_model = lpips.LPIPS(net='alex').to(device)
+
+# IQA metrics (BRISQUE, NIQE)
+# piq functions are stateless, import directly
+# Aesthetic predictor: use CLIP image-text ("a beautiful photograph") as proxy
+
+def extract_metadata(pil_img):
+    pm = ParserManager()
+    info = pm.parse(pil_img)
+    prompt = info.prompts[0].value if info.prompts else ''
+    model_name = info.model_name or ''
+    return prompt, model_name
+
+# Transform
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
+                         (0.26862954, 0.26130258, 0.27577711))
+])
+
+def compute_clip_score(img, text):
+    inputs = clip_processor(text=[text], images=img, return_tensors="pt", padding=True).to(device)
+    outputs = clip_model(**inputs)
+    score = torch.cosine_similarity(outputs.image_embeds, outputs.text_embeds)
+    return float((score.clamp(min=0) * 100).mean())
+
+@torch.no_grad()
+def compute_caption_similarity(img, prompt):
+    inputs = blip_processor(images=img, return_tensors="pt").to(device, torch.float16)
+    out = blip_model.generate(**inputs)
+    caption = blip_processor.decode(out[0], skip_special_tokens=True)
+    # similarity via CLIP
+    return compute_clip_score(img, caption)
+
+@torch.no_grad()
+def compute_iqa_metrics(img):
+    # convert to tensor
+    img_t = transforms.ToTensor()(img).unsqueeze(0).to(device)
+    brisque = float(piq.brisque(img_t).cpu())
+    niqe = float(piq.niqe(img_t).cpu())
+    return brisque, niqe
+
+@torch.no_grad()
+def compute_lpips_pair(img1, img2):
+    t1 = transforms.ToTensor()(img1).unsqueeze(0).to(device)
+    t2 = transforms.ToTensor()(img2).unsqueeze(0).to(device)
+    return float(lpips_model(t1, t2).cpu())
+
+# --------------------
+# Analysis Pipeline
+# --------------------
+
+def analyze_images(images):
+    records = []
+    imgs_by_model = {}
+
+    # extract metadata and compute per-image metrics
+    for img in images:
+        prompt, model = extract_metadata(img)
+        # basics
+        clip_score = compute_clip_score(img, prompt)
+        cap_sim = compute_caption_similarity(img, prompt)
+        brisque, niqe = compute_iqa_metrics(img)
+        # aesthetic proxy: CLIP with generic prompt
+        aest = compute_clip_score(img, "a beautiful high quality image")
+
+        records.append({
+            'model': model,
+            'prompt': prompt,
+            'clip_score': clip_score,
+            'caption_sim': cap_sim,
+            'brisque': brisque,
+            'niqe': niqe,
+            'aesthetic': aest,
+            'image': img
+        })
+        imgs_by_model.setdefault(model, []).append(img)
+
+    df = pd.DataFrame(records)
+
+    # diversity via LPIPS per model
+    diversity = {}
+    for model, imgs in imgs_by_model.items():
+        if len(imgs) < 2:
+            diversity[model] = 0.0
+        else:
+            pairs = []
+            for i in range(len(imgs)):
+                for j in range(i+1, len(imgs)):
+                    pairs.append(compute_lpips_pair(imgs[i], imgs[j]))
+            diversity[model] = sum(pairs) / len(pairs)
+
+    # aggregate
+    agg = df.groupby('model').agg({
+        'clip_score': ['mean'],
+        'caption_sim': ['mean'],
+        'brisque': ['mean'],
+        'niqe': ['mean'],
+        'aesthetic': ['mean']
+    })
+    agg.columns = ['_'.join(col) for col in agg.columns]
+    agg['diversity'] = pd.Series(diversity)
+    agg = agg.reset_index()
+
+    return df, agg
+
+# --------------------
+# Visualization Helpers
+# --------------------
+
+def plot_metrics(agg):
+    fig = px.bar(agg, x='model', y=['aesthetic_mean', 'clip_score_mean', 'caption_sim_mean', 'diversity'],
+                 barmode='group', title='Сравнение моделей по метрикам')
+    return fig
+
+# --------------------
+# Gradio Interface
+# --------------------
+
+def run_analysis(files):
+    images = [Image.open(io.BytesIO(f.read())).convert('RGB') for f in files]
+    df, agg = analyze_images(images)
+
+    # plots
+    fig = plot_metrics(agg)
+
+    return df, fig
+
+with gr.Blocks() as demo:
+    gr.Markdown("# AI Image Quality Evaluator")
+    gr.Markdown("Загрузите PNG-изображения, сгенерированные моделями ИИ, для анализа и сравнения моделей.")
+
+    with gr.Row():
+        input_files = gr.File(file_count="multiple", label="Выберите PNG файлы")
+        output_table = gr.DataFrame(headers=["model", "clip_score_mean", "caption_sim_mean", "brisque_mean", "niqe_mean", "aesthetic_mean", "diversity"], label="Сводная таблица")
+
+    plot_output = gr.Plot(label="График метрик")
+
+    run_btn = gr.Button("Запустить анализ")
+    run_btn.click(run_analysis, inputs=[input_files], outputs=[output_table, plot_output])
+
+if __name__ == "__main__":
+    demo.launch(server_name='0.0.0.0', share=False)