Scorer Starts

app.py

@@ -1,3 +1,119 @@
 import gradio as gr
 
-gr.load("models/q-future/one-align").launch()
+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 = ("""
+
+<h1 align="center">Q-Align: Teaching LMMs for Visual Scoring via Discrete Text-Defined Levels</h1>
+
+<h3 align="center"> One Unified Model for Visual scoring. </h3>
+
+<h5 align="center">
+      <a href="https://teowu.github.io/" target="_blank">Haoning Wu</a><sup>1</sup><sup>*</sup><sup>+</sup>,
+      <a href="https://github.com/zzc-1998" target="_blank">Zicheng Zhang</a><sup>2</sup><sup>*</sup>,
+    <a href="https://sites.google.com/view/r-panda" target="_blank">Weixia Zhang</a><sup>2</sup>,
+    <a href="https://chaofengc.github.io" target="_blank">Chaofeng Chen</a><sup>1</sup>,
+      <a href="https://liaoliang92.github.io" target="_blank">Liang Liao</a><sup>1</sup>,
+      <a href="https://github.com/lcysyzxdxc" target="_blank">Chunyi Li</a><sup>2</sup>,
+</h5>
+    
+
+<h5 align="center">
+        <a href="https://github.com/YixuanGao98" target="_blank">Yixuan Gao</a><sup>2</sup>,
+      <a href="https://github.com/AnnanWangDaniel" target="_blank">Annan Wang</a><sup>1</sup>,
+      <a href="https://github.com/ZhangErliCarl/" target="_blank">Erli Zhang</a><sup>1</sup>,
+      <a href="https://wenxiusun.com" target="_blank">Wenxiu Sun</a><sup>3</sup>,
+      <a href="https://scholar.google.com/citations?user=uT9CtPYAAAAJ&hl=en" target="_blank">Qiong Yan</a><sup>3</sup>,
+        <a href="https://sites.google.com/site/minxiongkuo/" target="_blank">Xiongkuo Min</a><sup>2</sup>,
+      <a href="https://ee.sjtu.edu.cn/en/FacultyDetail.aspx?id=24&infoid=153&flag=153" target="_blank">Guangtao Zhai</a><sup>2</sup><sup>#</sup>,
+      <a href="https://personal.ntu.edu.sg/wslin/Home.html" target="_blank">Weisi Lin</a><sup>1</sup><sup>#</sup>
+</h5>
+
+<h5 align="center">
+  <sup>1</sup>Nanyang Technological University, <sup>2</sup>Shanghai Jiao Tong University, <sup>3</sup>Sensetime Research
+</h5>
+<h5 align="center">
+<sup>*</sup>Equal contribution. <sup>+</sup>Project Lead. <sup>#</sup>Corresponding author(s). 
+</h5>
+
+<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)", info="If a video is uploaded, the image uploaded will be ignored.")
+
+labels = gr.Label(label="Probabilities of rating levels:")
+number = gr.Number(label="Output score:", info="Range in [1,5]. Higher is better.")
+demo = gr.Interface(fn=image_classifier, inputs=[input_img, input_vid, gr.Radio(["Image Aesthetics (IAA)", "Image Quality (IQA)", "Video Quality (VQA)"], label="Task", info="Which Scorer will you need?"),], outputs=[labels, number], title="OneScorer", description=title_markdown)
+demo.launch(share=True)

q_align/.ipynb_checkpoints/utils-checkpoint.py

@@ -0,0 +1,128 @@
+import datetime
+import logging
+import logging.handlers
+import os
+import sys
+
+import requests
+
+
+
+from q_align.constants import LOGDIR
+
+server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
+moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+
+handler = None
+
+
+def build_logger(logger_name, logger_filename):
+    global handler
+
+    formatter = logging.Formatter(
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    # Set the format of root handlers
+    if not logging.getLogger().handlers:
+        logging.basicConfig(level=logging.INFO)
+    logging.getLogger().handlers[0].setFormatter(formatter)
+
+    # Redirect stdout and stderr to loggers
+    stdout_logger = logging.getLogger("stdout")
+    stdout_logger.setLevel(logging.INFO)
+    sl = StreamToLogger(stdout_logger, logging.INFO)
+    sys.stdout = sl
+
+    stderr_logger = logging.getLogger("stderr")
+    stderr_logger.setLevel(logging.ERROR)
+    sl = StreamToLogger(stderr_logger, logging.ERROR)
+    sys.stderr = sl
+
+    # Get logger
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+
+    # Add a file handler for all loggers
+    if handler is None:
+        os.makedirs(LOGDIR, exist_ok=True)
+        filename = os.path.join(LOGDIR, logger_filename)
+        handler = logging.handlers.TimedRotatingFileHandler(
+            filename, when='D', utc=True)
+        handler.setFormatter(formatter)
+
+        for name, item in logging.root.manager.loggerDict.items():
+            if isinstance(item, logging.Logger):
+                item.addHandler(handler)
+
+    return logger
+
+
+class StreamToLogger(object):
+    """
+    Fake file-like stream object that redirects writes to a logger instance.
+    """
+    def __init__(self, logger, log_level=logging.INFO):
+        self.terminal = sys.stdout
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+
+    def __getattr__(self, attr):
+        return getattr(self.terminal, attr)
+
+    def write(self, buf):
+        temp_linebuf = self.linebuf + buf
+        self.linebuf = ''
+        for line in temp_linebuf.splitlines(True):
+            # From the io.TextIOWrapper docs:
+            #   On output, if newline is None, any '\n' characters written
+            #   are translated to the system default line separator.
+            # By default sys.stdout.write() expects '\n' newlines and then
+            # translates them so this is still cross platform.
+            if line[-1] == '\n':
+                self.logger.log(self.log_level, line.rstrip())
+            else:
+                self.linebuf += line
+
+    def flush(self):
+        if self.linebuf != '':
+            self.logger.log(self.log_level, self.linebuf.rstrip())
+        self.linebuf = ''
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def violates_moderation(text):
+    """
+    Check whether the text violates OpenAI moderation API.
+    """
+    url = "https://api.openai.com/v1/moderations"
+    headers = {"Content-Type": "application/json",
+               "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]}
+    text = text.replace("\n", "")
+    data = "{" + '"input": ' + f'"{text}"' + "}"
+    data = data.encode("utf-8")
+    try:
+        ret = requests.post(url, headers=headers, data=data, timeout=5)
+        flagged = ret.json()["results"][0]["flagged"]
+    except requests.exceptions.RequestException as e:
+        flagged = False
+    except KeyError as e:
+        flagged = False
+
+    return flagged
+
+
+def pretty_print_semaphore(semaphore):
+    if semaphore is None:
+        return "None"
+    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"

q_align/__init__.py

@@ -0,0 +1 @@
+from .model import MPLUGOwl2LlamaForCausalLM

q_align/constants.py

@@ -0,0 +1,9 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "./demo_logs"
+
+# Model Constants
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<|image|>"

q_align/conversation.py

@@ -0,0 +1,301 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+from q_align.constants import DEFAULT_IMAGE_TOKEN
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    TWO_NO_SYS = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            # init_msg = init_msg[0].replace("<image>", "").strip()
+            # if 'mmtag' in self.version:
+            #     messages[0] = (init_role, init_msg)
+            #     messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+            #     messages.insert(1, (self.roles[1], "Received."))
+            # else:
+            #     messages[0] = (init_role, "<image>\n" + init_msg)
+            init_msg = init_msg[0].replace(DEFAULT_IMAGE_TOKEN, "").strip()
+            messages[0] = (init_role, DEFAULT_IMAGE_TOKEN + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO_NO_SYS:
+            seps = [self.sep, self.sep2]
+            ret = ""
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    msg, image, image_process_mode = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<|image|>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_mplug_owl2 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO_NO_SYS,
+    sep=" ",
+    sep2="</s>",
+)
+
+# default_conversation = conv_vicuna_v1
+default_conversation = conv_mplug_owl2
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "mplug_owl2": conv_mplug_owl2,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

q_align/evaluate/.ipynb_checkpoints/iaa_eval-checkpoint.py

@@ -0,0 +1,164 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+from scipy.stats import spearmanr, pearsonr
+
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+def wa5(logits):
+    import numpy as np
+    logprobs = np.array([logits["excellent"], logits["good"], logits["fair"], logits["poor"], logits["bad"]])
+    probs = np.exp(logprobs) / np.sum(np.exp(logprobs))
+    return np.inner(probs, np.array([1,0.75,0.5,0.25,0.]))
+
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_path = "playground/data/"
+                  
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        json_prefix + "test_ava.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the aesthetics of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp = DEFAULT_IMAGE_TOKEN + inp
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The aesthetics of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for json_ in jsons:
+        with open(json_) as f:
+            iqadata = json.load(f)  
+
+            image_tensors = []
+            batch_data = []
+            prs, gts = [], []
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                filename = llddata["image"]
+                llddata["logits"] = defaultdict(float)
+                
+                
+                
+                image = load_image(image_path + filename)
+                def expand2square(pil_img, background_color):
+                        width, height = pil_img.size
+                        if width == height:
+                            return pil_img
+                        elif width > height:
+                            result = Image.new(pil_img.mode, (width, width), background_color)
+                            result.paste(pil_img, (0, (width - height) // 2))
+                            return result
+                        else:
+                            result = Image.new(pil_img.mode, (height, height), background_color)
+                            result.paste(pil_img, ((height - width) // 2, 0))
+                            return result
+                image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+                image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+                image_tensors.append(image_tensor)
+                batch_data.append(llddata)
+
+                if i % 8 == 7 or i == len(iqadata) - 1:                     
+                    with torch.inference_mode():
+                        output_logits = model(input_ids.repeat(len(image_tensors), 1),
+                            images=torch.cat(image_tensors, 0))["logits"][:,-1]
+    
+                    for j, xllddata in enumerate(batch_data):
+                        for tok, id_ in zip(toks, ids_):
+                            xllddata["logits"][tok] += output_logits[j,id_].item()
+                        xllddata["score"] = wa5(xllddata["logits"])
+                        # print(llddata)
+                        prs.append(xllddata["score"])
+                        gts.append(xllddata["gt_score"])
+                        json_ = json_.replace("combined/", "combined-")
+                        with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
+                            json.dump(xllddata, wf)
+
+                    image_tensors = []
+                    batch_data = []
+                
+                #if i > 0 and i % 200 == 0:
+                #    print(spearmanr(prs,gts)[0], pearsonr(prs,gts)[0])
+            print("Spearmanr", spearmanr(prs,gts)[0], "Pearson", pearsonr(prs,gts)[0])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/.ipynb_checkpoints/iqa4vqa_eval-checkpoint.py

@@ -0,0 +1,150 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+from decord import VideoReader
+
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_video(video_file):
+    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))]
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_paths = [
+        "playground/data/",
+        "playground/data/",
+        "playground/data/KoNViD_1k_videos/",
+        "playground/data/maxwell/",
+    ]
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        json_prefix + "test_lsvq.json",
+        json_prefix + "test_lsvq_1080p.json",
+        json_prefix + "konvid.json",
+        json_prefix + "maxwell_test.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the quality of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for image_path, json_ in zip(image_paths, jsons):
+        with open(json_) as f:
+            iqadata = json.load(f) 
+            try:
+                for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                    filename = llddata["img_path"]
+                    llddata["logits"] = defaultdict(float)
+
+                    image = load_video(image_path + filename)
+                    def expand2square(pil_img, background_color):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                    image = [expand2square(img, tuple(int(x*255) for x in image_processor.image_mean)) for img in image]
+                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+
+                    if True:
+                        with torch.inference_mode():
+                            output_logits = model(input_ids.repeat(image_tensor.shape[0], 1),
+                                images=image_tensor)["logits"][:,-1]
+
+                            for tok, id_ in zip(toks, ids_):
+                                llddata["logits"][tok] += output_logits.mean(0)[id_].item()
+                            # print(llddata)
+                            json_ = json_.replace("combined/", "combined-")
+                            with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
+                                json.dump(llddata, wf)
+            except:
+                continue
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/q-align-image")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/.ipynb_checkpoints/iqa_eval-checkpoint.py

@@ -0,0 +1,156 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_path = "playground/data/"
+                  
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        json_prefix + "test_imagerewarddb.json",
+        json_prefix + "test_koniq.json",
+        json_prefix + "test_spaq.json",
+        json_prefix + "test_kadid.json",
+        json_prefix + "livec.json",
+        json_prefix + "agi.json",
+        json_prefix + "live.json",
+        json_prefix + "csiq.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "Evaluate the image quality of the following image."#"How would you rate the quality of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for json_ in jsons:
+        with open(json_) as f:
+            iqadata = json.load(f)  
+
+            image_tensors = []
+            batch_data = []
+            
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                if True:
+                    try:
+                        filename = llddata["image"]
+                    except:
+                        filename = llddata["img_path"]
+                    llddata["logits"] = defaultdict(float)
+
+                    image = load_image(image_path + filename)
+                    def expand2square(pil_img, background_color):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                    image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+                    image_tensors.append(image_tensor)
+                    batch_data.append(llddata)
+
+                    if i % 8 == 7 or i == len(iqadata) - 1:                     
+                        with torch.inference_mode():
+                            output_logits = model(input_ids.repeat(len(image_tensors), 1),
+                                images=torch.cat(image_tensors, 0))["logits"][:,-1]
+
+                        for j, xllddata in enumerate(batch_data):
+                            for tok, id_ in zip(toks, ids_):
+                                xllddata["logits"][tok] += output_logits[j,id_].item()
+                            # print(llddata)
+                            json_ = json_.replace("combined/", "combined-")
+                            with open(f"results/{args.model_path}/2{json_.split('/')[-1]}", "a") as wf:
+                                json.dump(xllddata, wf)
+
+                        image_tensors = []
+                        batch_data = []
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/.ipynb_checkpoints/scorer-checkpoint.py

@@ -0,0 +1,155 @@
+from PIL import Image
+
+import torch.nn as nn
+import torch
+
+from typing import List
+
+from q_align.model.builder import load_pretrained_model
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+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))]
+
+
+class QAlignScorer(nn.Module):
+    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
+        super().__init__()
+        if model is None:
+            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
+        prompt = "USER: How would you rate the quality of this image?\n<|image|>\nASSISTANT: The quality of the image is"
+        
+        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
+        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
+    
+        self.tokenizer = tokenizer
+        self.model = model
+        self.image_processor = image_processor
+        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+        
+    def forward(self, image: List[Image.Image]):
+        image = [self.expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in image]
+        with torch.inference_mode():
+            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"].half().to(self.model.device)
+            output_logits = self.model(self.input_ids.repeat(image_tensor.shape[0], 1),
+                        images=image_tensor)["logits"][:,-1, self.preferential_ids_]
+
+            return torch.softmax(output_logits, -1) #@ self.weight_tensor
+        
+        
+class QAlignAestheticScorer(nn.Module):
+    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
+        super().__init__()
+        if model is None:
+            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
+        prompt = "USER: How would you rate the aesthetics of this image?\n<|image|>\nASSISTANT: The aesthetics of the image is"
+        
+        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
+        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
+    
+        self.tokenizer = tokenizer
+        self.model = model
+        self.image_processor = image_processor
+        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+        
+    def forward(self, image: List[Image.Image]):
+        image = [self.expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in image]
+        with torch.inference_mode():
+            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"].half().to(self.model.device)
+            output_logits = self.model(self.input_ids.repeat(image_tensor.shape[0], 1),
+                        images=image_tensor)["logits"][:,-1, self.preferential_ids_]
+
+            return torch.softmax(output_logits, -1) #@ self.weight_tensor
+        
+class QAlignVideoScorer(nn.Module):
+    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
+        super().__init__()
+        if model is None:
+            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
+        prompt = "USER: How would you rate the quality of this video?\n<|image|>\nASSISTANT: The quality of the video is"
+        
+        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
+        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
+    
+        self.tokenizer = tokenizer
+        self.model = model
+        self.image_processor = image_processor
+        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+        
+    def forward(self, video: List[List[Image.Image]]):
+        video = [[self.expand2square(frame, tuple(int(x*255) for x in self.image_processor.image_mean)) for frame in vid] for vid in video]
+        with torch.inference_mode():
+            video_tensors = [self.image_processor.preprocess(vid, return_tensors="pt")["pixel_values"].half().to(self.model.device) for vid in video]
+            output_logits = self.model(self.input_ids.repeat(len(video_tensors), 1),
+                        images=video_tensors)["logits"][:,-1, self.preferential_ids_]
+            return torch.softmax(output_logits, -1) #@ self.weight_tensor
+    
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--img_path", type=str, default="fig/singapore_flyer.jpg")
+    parser.add_argument("--aesthetic", action="store_true")
+    parser.add_argument("--video", action="store_true")
+    args = parser.parse_args()
+
+    if args.video:
+        scorer = QAlignVideoScorer(pretrained=args.model_path, device=args.device)
+        print(scorer([load_video(args.img_path)]).tolist())
+    else:
+        scorer = QAlignScorer(pretrained=args.model_path, device=args.device) if not args.aesthetic else QAlignAestheticScorer(pretrained=args.model_path, device=args.device) 
+        print(scorer([Image.open(args.img_path)]).tolist())
+

q_align/evaluate/.ipynb_checkpoints/vqa_eval-checkpoint.py

@@ -0,0 +1,167 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+
+from scipy.stats import spearmanr, pearsonr
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+def wa5(logits):
+    import numpy as np
+    logprobs = np.array([logits["excellent"], logits["good"], logits["fair"], logits["poor"], logits["bad"]])
+    probs = np.exp(logprobs) / np.sum(np.exp(logprobs))
+    return np.inner(probs, np.array([1,0.75,0.5,0.25,0.]))
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+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))]
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_paths = [
+        #"playground/data/",
+        #"playground/data/",
+        "playground/data/KoNViD_1k_videos/",
+        "playground/data/maxwell/",
+
+    ]
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        #json_prefix + "test_lsvq.json",
+        #json_prefix + "test_lsvq_1080p.json",
+        json_prefix + "konvid.json",
+        json_prefix + "maxwell_test.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the quality of this video?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the video is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for image_path, json_ in zip(image_paths, jsons):
+        with open(json_) as f:
+            iqadata = json.load(f) 
+            prs, gts = [], []
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                try:
+                    try:
+                        filename = llddata["img_path"]
+                    except:
+                        filename = llddata["image"]
+                    llddata["logits"] = defaultdict(float)
+
+                    image = load_video(image_path + filename)
+                    def expand2square(pil_img, background_color):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                    image = [expand2square(img, tuple(int(x*255) for x in image_processor.image_mean)) for img in image]
+                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+                    if True:
+                        with torch.inference_mode():
+                            output_logits = model(input_ids,
+                                images=[image_tensor])["logits"][:,-1]
+                            for tok, id_ in zip(toks, ids_):
+                                llddata["logits"][tok] += output_logits.mean(0)[id_].item()
+                            llddata["score"] = wa5(llddata["logits"])
+                            # print(llddata)
+                            prs.append(llddata["score"])
+                            gts.append(llddata["gt_score"])
+                            # print(llddata)
+                            json_ = json_.replace("combined/", "combined-")
+                            with open(f"results/{args.model_path}/2{json_.split('/')[-1]}", "a") as wf:
+                                json.dump(llddata, wf)
+
+                    if i > 0 and i % 200 == 0:
+                        print(spearmanr(prs,gts)[0], pearsonr(prs,gts)[0])
+                except:
+                    continue
+            print("Spearmanr", spearmanr(prs,gts)[0], "Pearson", pearsonr(prs,gts)[0])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/eval.py

@@ -0,0 +1,138 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_path = "playground/data/"
+
+    json_prefix = "playground/data/labels/mos_simple/"
+    jsons = [
+        json_prefix + "test_flive.json",
+        json_prefix + "combined/kadid_ref.json",
+        json_prefix + "combined/livec.json",
+        json_prefix + "test_koniq.json",
+        json_prefix + "test_spaq.json",
+        json_prefix + "combined/agi.json",
+        json_prefix + "combined/kadid.json",   
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the quality of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+    
+    for json_ in jsons:
+        with open(json_) as f:
+            iqadata = json.load(f)  
+
+            image_tensors = []
+            batch_data = []
+            
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                #print(f"Evaluating image {i}")
+                #print(prompt)
+                filename = llddata["image"]
+                llddata["logits"] = defaultdict(float)
+                
+                image = load_image(image_path + filename)
+                image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().cuda()
+
+                image_tensors.append(image_tensor)
+                batch_data.append(llddata)
+
+                if i % 8 == 7 or i == len(iqadata) - 1:                     
+                    with torch.inference_mode():
+                        output_logits = model(input_ids.repeat(len(image_tensors), 1),
+                            images=torch.cat(image_tensors, 0))["logits"][:,-1]
+    
+                    for j, xllddata in enumerate(batch_data):
+                        for tok, id_ in zip(toks, ids_):
+                            xllddata["logits"][tok] += output_logits[j,id_].item()
+                        # print(llddata)
+                        json_ = json_.replace("combined/", "combined-")
+                        # print(f"results/mix-mplug-owl-2-boost_iqa_wu_v2/{json_.split('/')[-1]}")    
+                        with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
+                            json.dump(xllddata, wf)
+
+                    image_tensors = []
+                    batch_data = []
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/q-align-koniq-spaq-v0")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/iaa_eval.py

@@ -0,0 +1,164 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+from scipy.stats import spearmanr, pearsonr
+
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+def wa5(logits):
+    import numpy as np
+    logprobs = np.array([logits["excellent"], logits["good"], logits["fair"], logits["poor"], logits["bad"]])
+    probs = np.exp(logprobs) / np.sum(np.exp(logprobs))
+    return np.inner(probs, np.array([1,0.75,0.5,0.25,0.]))
+
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_path = "playground/data/"
+                  
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        json_prefix + "test_ava.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the aesthetics of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp = DEFAULT_IMAGE_TOKEN + inp
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The aesthetics of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for json_ in jsons:
+        with open(json_) as f:
+            iqadata = json.load(f)  
+
+            image_tensors = []
+            batch_data = []
+            prs, gts = [], []
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                filename = llddata["image"]
+                llddata["logits"] = defaultdict(float)
+                
+                
+                
+                image = load_image(image_path + filename)
+                def expand2square(pil_img, background_color):
+                        width, height = pil_img.size
+                        if width == height:
+                            return pil_img
+                        elif width > height:
+                            result = Image.new(pil_img.mode, (width, width), background_color)
+                            result.paste(pil_img, (0, (width - height) // 2))
+                            return result
+                        else:
+                            result = Image.new(pil_img.mode, (height, height), background_color)
+                            result.paste(pil_img, ((height - width) // 2, 0))
+                            return result
+                image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+                image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+                image_tensors.append(image_tensor)
+                batch_data.append(llddata)
+
+                if i % 8 == 7 or i == len(iqadata) - 1:                     
+                    with torch.inference_mode():
+                        output_logits = model(input_ids.repeat(len(image_tensors), 1),
+                            images=torch.cat(image_tensors, 0))["logits"][:,-1]
+    
+                    for j, xllddata in enumerate(batch_data):
+                        for tok, id_ in zip(toks, ids_):
+                            xllddata["logits"][tok] += output_logits[j,id_].item()
+                        xllddata["score"] = wa5(xllddata["logits"])
+                        # print(llddata)
+                        prs.append(xllddata["score"])
+                        gts.append(xllddata["gt_score"])
+                        json_ = json_.replace("combined/", "combined-")
+                        with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
+                            json.dump(xllddata, wf)
+
+                    image_tensors = []
+                    batch_data = []
+                
+                #if i > 0 and i % 200 == 0:
+                #    print(spearmanr(prs,gts)[0], pearsonr(prs,gts)[0])
+            print("Spearmanr", spearmanr(prs,gts)[0], "Pearson", pearsonr(prs,gts)[0])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/iqa4vqa_eval.py

@@ -0,0 +1,150 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+from decord import VideoReader
+
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_video(video_file):
+    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))]
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_paths = [
+        "playground/data/",
+        "playground/data/",
+        "playground/data/KoNViD_1k_videos/",
+        "playground/data/maxwell/",
+    ]
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        json_prefix + "test_lsvq.json",
+        json_prefix + "test_lsvq_1080p.json",
+        json_prefix + "konvid.json",
+        json_prefix + "maxwell_test.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the quality of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for image_path, json_ in zip(image_paths, jsons):
+        with open(json_) as f:
+            iqadata = json.load(f) 
+            try:
+                for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                    filename = llddata["img_path"]
+                    llddata["logits"] = defaultdict(float)
+
+                    image = load_video(image_path + filename)
+                    def expand2square(pil_img, background_color):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                    image = [expand2square(img, tuple(int(x*255) for x in image_processor.image_mean)) for img in image]
+                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+
+                    if True:
+                        with torch.inference_mode():
+                            output_logits = model(input_ids.repeat(image_tensor.shape[0], 1),
+                                images=image_tensor)["logits"][:,-1]
+
+                            for tok, id_ in zip(toks, ids_):
+                                llddata["logits"][tok] += output_logits.mean(0)[id_].item()
+                            # print(llddata)
+                            json_ = json_.replace("combined/", "combined-")
+                            with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
+                                json.dump(llddata, wf)
+            except:
+                continue
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/q-align-image")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/iqa_eval.py

@@ -0,0 +1,156 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_path = "playground/data/"
+                  
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        json_prefix + "test_imagerewarddb.json",
+        json_prefix + "test_koniq.json",
+        json_prefix + "test_spaq.json",
+        json_prefix + "test_kadid.json",
+        json_prefix + "livec.json",
+        json_prefix + "agi.json",
+        json_prefix + "live.json",
+        json_prefix + "csiq.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "Evaluate the image quality of the following image."#"How would you rate the quality of this image?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the image is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for json_ in jsons:
+        with open(json_) as f:
+            iqadata = json.load(f)  
+
+            image_tensors = []
+            batch_data = []
+            
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                if True:
+                    try:
+                        filename = llddata["image"]
+                    except:
+                        filename = llddata["img_path"]
+                    llddata["logits"] = defaultdict(float)
+
+                    image = load_image(image_path + filename)
+                    def expand2square(pil_img, background_color):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                    image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+                    image_tensors.append(image_tensor)
+                    batch_data.append(llddata)
+
+                    if i % 8 == 7 or i == len(iqadata) - 1:                     
+                        with torch.inference_mode():
+                            output_logits = model(input_ids.repeat(len(image_tensors), 1),
+                                images=torch.cat(image_tensors, 0))["logits"][:,-1]
+
+                        for j, xllddata in enumerate(batch_data):
+                            for tok, id_ in zip(toks, ids_):
+                                xllddata["logits"][tok] += output_logits[j,id_].item()
+                            # print(llddata)
+                            json_ = json_.replace("combined/", "combined-")
+                            with open(f"results/{args.model_path}/2{json_.split('/')[-1]}", "a") as wf:
+                                json.dump(xllddata, wf)
+
+                        image_tensors = []
+                        batch_data = []
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/evaluate/scorer.py

@@ -0,0 +1,155 @@
+from PIL import Image
+
+import torch.nn as nn
+import torch
+
+from typing import List
+
+from q_align.model.builder import load_pretrained_model
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+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))]
+
+
+class QAlignScorer(nn.Module):
+    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
+        super().__init__()
+        if model is None:
+            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
+        prompt = "USER: How would you rate the quality of this image?\n<|image|>\nASSISTANT: The quality of the image is"
+        
+        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
+        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
+    
+        self.tokenizer = tokenizer
+        self.model = model
+        self.image_processor = image_processor
+        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+        
+    def forward(self, image: List[Image.Image]):
+        image = [self.expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in image]
+        with torch.inference_mode():
+            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"].half().to(self.model.device)
+            output_logits = self.model(self.input_ids.repeat(image_tensor.shape[0], 1),
+                        images=image_tensor)["logits"][:,-1, self.preferential_ids_]
+
+            return torch.softmax(output_logits, -1) #@ self.weight_tensor
+        
+        
+class QAlignAestheticScorer(nn.Module):
+    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
+        super().__init__()
+        if model is None:
+            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
+        prompt = "USER: How would you rate the aesthetics of this image?\n<|image|>\nASSISTANT: The aesthetics of the image is"
+        
+        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
+        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
+    
+        self.tokenizer = tokenizer
+        self.model = model
+        self.image_processor = image_processor
+        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+        
+    def forward(self, image: List[Image.Image]):
+        image = [self.expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in image]
+        with torch.inference_mode():
+            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"].half().to(self.model.device)
+            output_logits = self.model(self.input_ids.repeat(image_tensor.shape[0], 1),
+                        images=image_tensor)["logits"][:,-1, self.preferential_ids_]
+
+            return torch.softmax(output_logits, -1) #@ self.weight_tensor
+        
+class QAlignVideoScorer(nn.Module):
+    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
+        super().__init__()
+        if model is None:
+            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
+        prompt = "USER: How would you rate the quality of this video?\n<|image|>\nASSISTANT: The quality of the video is"
+        
+        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
+        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
+    
+        self.tokenizer = tokenizer
+        self.model = model
+        self.image_processor = image_processor
+        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+        
+    def forward(self, video: List[List[Image.Image]]):
+        video = [[self.expand2square(frame, tuple(int(x*255) for x in self.image_processor.image_mean)) for frame in vid] for vid in video]
+        with torch.inference_mode():
+            video_tensors = [self.image_processor.preprocess(vid, return_tensors="pt")["pixel_values"].half().to(self.model.device) for vid in video]
+            output_logits = self.model(self.input_ids.repeat(len(video_tensors), 1),
+                        images=video_tensors)["logits"][:,-1, self.preferential_ids_]
+            return torch.softmax(output_logits, -1) #@ self.weight_tensor
+    
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--img_path", type=str, default="fig/singapore_flyer.jpg")
+    parser.add_argument("--aesthetic", action="store_true")
+    parser.add_argument("--video", action="store_true")
+    args = parser.parse_args()
+
+    if args.video:
+        scorer = QAlignVideoScorer(pretrained=args.model_path, device=args.device)
+        print(scorer([load_video(args.img_path)]).tolist())
+    else:
+        scorer = QAlignScorer(pretrained=args.model_path, device=args.device) if not args.aesthetic else QAlignAestheticScorer(pretrained=args.model_path, device=args.device) 
+        print(scorer([Image.open(args.img_path)]).tolist())
+

q_align/evaluate/vqa_eval.py

@@ -0,0 +1,167 @@
+import argparse
+import torch
+
+from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from q_align.conversation import conv_templates, SeparatorStyle
+from q_align.model.builder import load_pretrained_model
+from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+
+from scipy.stats import spearmanr, pearsonr
+
+import json
+from tqdm import tqdm
+from collections import defaultdict
+
+import os
+
+def wa5(logits):
+    import numpy as np
+    logprobs = np.array([logits["excellent"], logits["good"], logits["fair"], logits["poor"], logits["bad"]])
+    probs = np.exp(logprobs) / np.sum(np.exp(logprobs))
+    return np.inner(probs, np.array([1,0.75,0.5,0.25,0.]))
+
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+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))]
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+    
+    
+    import json
+
+    
+    image_paths = [
+        #"playground/data/",
+        #"playground/data/",
+        "playground/data/KoNViD_1k_videos/",
+        "playground/data/maxwell/",
+
+    ]
+
+    json_prefix = "playground/data/test_jsons/"
+    jsons = [
+        #json_prefix + "test_lsvq.json",
+        #json_prefix + "test_lsvq_1080p.json",
+        json_prefix + "konvid.json",
+        json_prefix + "maxwell_test.json",
+    ]
+
+    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
+
+
+    conv_mode = "mplug_owl2"
+    
+    inp = "How would you rate the quality of this video?"
+        
+    conv = conv_templates[conv_mode].copy()
+    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
+    conv.append_message(conv.roles[0], inp)
+    image = None
+        
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt() + " The quality of the video is"
+    
+    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
+    print(toks)
+    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
+    print(ids_)
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
+    
+    for image_path, json_ in zip(image_paths, jsons):
+        with open(json_) as f:
+            iqadata = json.load(f) 
+            prs, gts = [], []
+            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
+                try:
+                    try:
+                        filename = llddata["img_path"]
+                    except:
+                        filename = llddata["image"]
+                    llddata["logits"] = defaultdict(float)
+
+                    image = load_video(image_path + filename)
+                    def expand2square(pil_img, background_color):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                    image = [expand2square(img, tuple(int(x*255) for x in image_processor.image_mean)) for img in image]
+                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
+
+                    if True:
+                        with torch.inference_mode():
+                            output_logits = model(input_ids,
+                                images=[image_tensor])["logits"][:,-1]
+                            for tok, id_ in zip(toks, ids_):
+                                llddata["logits"][tok] += output_logits.mean(0)[id_].item()
+                            llddata["score"] = wa5(llddata["logits"])
+                            # print(llddata)
+                            prs.append(llddata["score"])
+                            gts.append(llddata["gt_score"])
+                            # print(llddata)
+                            json_ = json_.replace("combined/", "combined-")
+                            with open(f"results/{args.model_path}/2{json_.split('/')[-1]}", "a") as wf:
+                                json.dump(llddata, wf)
+
+                    if i > 0 and i % 200 == 0:
+                        print(spearmanr(prs,gts)[0], pearsonr(prs,gts)[0])
+                except:
+                    continue
+            print("Spearmanr", spearmanr(prs,gts)[0], "Pearson", pearsonr(prs,gts)[0])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="q-future/one-align")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    args = parser.parse_args()
+    main(args)

q_align/mm_utils.py

@@ -0,0 +1,112 @@
+from PIL import Image
+from io import BytesIO
+import base64
+
+import torch
+from transformers import StoppingCriteria
+from q_align.constants import IMAGE_TOKEN_INDEX,DEFAULT_IMAGE_TOKEN
+from icecream import ic
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def process_images(images, image_processor, model_cfg=None):
+    if model_cfg is not None:
+        image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+    else:
+        image_aspect_ratio = 'resize'
+    new_images = []
+    if image_aspect_ratio == 'pad':
+        for image in images:
+            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            new_images.append(image)
+    elif image_aspect_ratio == 'resize':
+        for image in images:
+            max_edge = max(image.size)
+            image = image.resize((max_edge, max_edge))
+            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            new_images.append(image)
+    else:
+        return image_processor(images, return_tensors='pt')['pixel_values']
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
+
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    prompt_chunks = [tokenizer(chunk).input_ids if len(chunk) > 0 else [] for chunk in prompt.split(DEFAULT_IMAGE_TOKEN)]
+
+    def insert_separator(X, sep):
+        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+
+    input_ids = []
+    offset = 0
+    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+        offset = 1
+        input_ids.append(prompt_chunks[0][0])
+
+    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+        input_ids.extend(x[offset:])
+
+    if return_tensors is not None:
+        if return_tensors == 'pt':
+            return torch.tensor(input_ids, dtype=torch.long)
+        raise ValueError(f'Unsupported tensor type: {return_tensors}')
+    return input_ids
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith('checkpoint-'):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)"  # TODO
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            if (output_ids[0, -keyword_id.shape[0]:] == keyword_id).all():
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False

q_align/model/__init__.py

@@ -0,0 +1,2 @@
+from .modeling_mplug_owl2 import MPLUGOwl2LlamaForCausalLM
+from .configuration_mplug_owl2 import MPLUGOwl2Config