Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import argparse | |
| import datetime | |
| import json | |
| import os | |
| import time | |
| import gradio as gr | |
| import requests | |
| from PIL import Image | |
| from q_align.model.builder import load_pretrained_model | |
| from q_align.conversation import (default_conversation, conv_templates, | |
| SeparatorStyle) | |
| from q_align.constants import LOGDIR | |
| from q_align.utils import (build_logger, server_error_msg, | |
| violates_moderation, moderation_msg) | |
| from q_align.evaluate.scorer import QAlignScorer, QAlignAestheticScorer, QAlignVideoScorer | |
| import gradio as gr | |
| def load_video(video_file): | |
| from decord import VideoReader | |
| vr = VideoReader(video_file) | |
| # Get video frame rate | |
| fps = vr.get_avg_fps() | |
| # Calculate frame indices for 1fps | |
| frame_indices = [int(fps * i) for i in range(int(len(vr) / fps))] | |
| frames = vr.get_batch(frame_indices).asnumpy() | |
| return [Image.fromarray(frames[i]) for i in range(int(len(vr) / fps))] | |
| pretrained="q-future/one-align" | |
| device="cuda:0" | |
| tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device) | |
| iqa_scorer = QAlignScorer(tokenizer=tokenizer, model=model, image_processor=image_processor) | |
| iaa_scorer = QAlignAestheticScorer(tokenizer=tokenizer, model=model, image_processor=image_processor) | |
| vqa_scorer = QAlignVideoScorer(tokenizer=tokenizer, model=model, image_processor=image_processor) | |
| scorers = {"Image Aesthetics (IAA)": iaa_scorer, "Image Quality (IQA)": iqa_scorer, "Video Quality (VQA)": vqa_scorer} | |
| LEVELS = ["excellent (5)", "good (4)", "fair (3)", "poor (2)", "bad (1)"] | |
| scores = [5,4,3,2,1] | |
| def image_classifier(input_img, input_vid, scorer_type): | |
| if scorer_type is None: | |
| scorer_type = "Image Quality (IQA)" | |
| this_scorer = scorers[scorer_type] | |
| if input_vid is not None: | |
| input_ = load_video(input_vid) | |
| elif input_img is not None: | |
| input_ = [input_img] | |
| if "Video" in scorer_type: | |
| input_ = [input_] | |
| probs = this_scorer(input_).mean(0).tolist() | |
| prob_dict = {LEVEL: prob for LEVEL, prob in zip(LEVELS, probs)} | |
| score = sum([prob * score for score, prob in zip(scores, probs)]) | |
| return prob_dict, score | |
| title_markdown = (""" | |
| <div style="width: 100%; text-align: center; margin:auto;"> | |
| <img style="width: 100%" src="https://raw.githubusercontent.com/Q-Future/Q-Align/main/fig/onescorer.png"> | |
| </div> | |
| <h4 align="center"> If you like the OneScorer, please give us a star ✨ on <a href='https://github.com/Q-Future/Q-Align'>[GitHub]</a> for latest update. </h4> | |
| <h5 align="center"> | |
| <div style="display:flex; gap: 0.25rem;" align="center"> | |
| <a href='https://q-align.github.io'><img src='https://img.shields.io/badge/Homepage-green'></a> | |
| <a href='https://github.com/Q-Future/Q-Align'><img src='https://img.shields.io/badge/Github-Code-blue'></a> | |
| <a href="https://Q-Future.github.io/Q-Align/fig/Q_Align_v0_1_preview.pdf"><img src="https://img.shields.io/badge/Technical-Report-red"></a> | |
| <a href='https://github.com/Q-Future/Q-Align/stargazers'><img src='https://img.shields.io/github/stars/Q-Future/Q-Align.svg?style=social'></a> | |
| </div> | |
| </h5> | |
| """) | |
| input_img = gr.Image(type='pil', label="Upload an Image") | |
| input_vid = gr.Video(label="Upload a Video (will INGORE the image if a video is uploaded)",sources=["upload"]) | |
| radio = gr.Radio(["Image Aesthetics (IAA)", "Image Quality (IQA)", "Video Quality (VQA)"], label="Task", info="Which Scorer will you need?") | |
| input_img = gr.Image(type='pil', label="Upload an Image") | |
| labels = gr.Label(label="Probabilities of rating levels:") | |
| number = gr.Number(label="Output score:", info="Range in [1,5]. Higher is better.", precision=4) | |
| demo = gr.Interface(fn=image_classifier, inputs=[input_img, input_vid, radio], outputs=[labels, number], description=title_markdown, examples=[["fig/eiffel_a.jpg", None, "Image Aesthetics (IAA)"], ["fig/singapore_flyer_2.jpg", None, "Image Quality (IQA)"], ["fig/none.png", "fig/10244479353.mp4", "Video Quality (VQA)"]], article="This is the Scorer Demo as Proposed by Paper: 'Q-Align: Teaching LMMs for Visual Scoring via Discrete Text-Defined Levels'. The proposed Q-Align achieves state-of-the-art performance on image quality assessment (IQA), image aesthetic assessment (IAA), as well as video quality assessment (VQA) tasks under the original LMM structure. With the syllabus, we further unify the three tasks into one model, termed the **OneAlign**, to which the demo corresponds.") | |
| demo.launch() | |