Views
No views yet
The first NR-IQA model enhanced by RL2R, capable of both quality description and rating through reasoning.

PROMPT = (
"You are doing the image quality assessment task. Here is the question: "
"What is your overall rating on the quality of this picture? The rating should be a float between 1 and 5, "
"rounded to two decimal places, with 1 representing very poor quality and 5 representing excellent quality."
)
QUESTION_TEMPLATE = "{Question} Please only output the final answer with only one score in <answer> </answer> tags."1from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
2from qwen_vl_utils import process_vision_info
3
4import torch
5import random
6import re
7import os
8
9
10def score_image(image_path, model, processor):
11 PROMPT = (
12 "You are doing the image quality assessment task. Here is the question: "
13 "What is your overall rating on the quality of this picture? The rating should be a float between 1 and 5, "
14 "rounded to two decimal places, with 1 representing very poor quality and 5 representing excellent quality."
15 )
16
17 QUESTION_TEMPLATE = "{Question} Please only output the final answer with only one score in <answer> </answer> tags."
18 message = [
19 {
20 "role": "user",
21 "content": [
22 {'type': 'image', 'image': image_path},
23 {"type": "text", "text": QUESTION_TEMPLATE.format(Question=PROMPT)}
24 ],
25 }
26 ]
27
28 batch_messages = [message]
29
30 # Preparation for inference
31 text = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True, add_vision_id=True) for msg in batch_messages]
32 image_inputs, video_inputs = process_vision_info(batch_messages)
33 inputs = processor(
34 text=text,
35 images=image_inputs,
36 videos=video_inputs,
37 padding=True,
38 return_tensors="pt",
39 )
40 inputs = inputs.to(device)
41
42 # Inference: Generation of the output
43 generated_ids = model.generate(**inputs, use_cache=True, max_new_tokens=2048, do_sample=True, top_k=50, top_p=1)
44 generated_ids_trimmed = [
45 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
46 ]
47 batch_output_text = processor.batch_decode(
48 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
49 )
50
51 reasoning = None
52
53 try:
54 model_output_matches = re.findall(r'<answer>(.*?)</answer>', batch_output_text[0], re.DOTALL)
55 model_answer = model_output_matches[-1].strip() if model_output_matches else batch_output_text[0].strip()
56 score = float(re.search(r'\d+(\.\d+)?', model_answer).group())
57 except:
58 print(f"================= Meet error with {img_path}, please generate again. =================")
59 score = random.randint(1, 5)
60
61 return reasoning, score
62
63
64random.seed(1)
65MODEL_PATH = ""
66device = torch.device("cuda:5") if torch.cuda.is_available() else torch.device("cpu")
67image_path = ""
68
69model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
70 MODEL_PATH,
71 torch_dtype=torch.bfloat16,
72 attn_implementation="flash_attention_2",
73 device_map=device,
74)
75processor = AutoProcessor.from_pretrained(MODEL_PATH)
76processor.tokenizer.padding_side = "left"
77
78reasoning, score = score_image(
79 image_path, model, processor
80)
81
82print(score)1from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
2from qwen_vl_utils import process_vision_info
3from tqdm import tqdm
4
5import torch
6import random
7import re
8import os
9
10
11def get_image_paths(folder_path):
12 image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}
13 image_paths = []
14
15 for root, dirs, files in os.walk(folder_path):
16 for file in files:
17 _, ext = os.path.splitext(file)
18 if ext.lower() in image_extensions:
19 image_paths.append(os.path.join(root, file))
20
21 return image_paths
22
23def score_batch_image(image_paths, model, processor):
24 PROMPT = (
25 "You are doing the image quality assessment task. Here is the question: "
26 "What is your overall rating on the quality of this picture? The rating should be a float between 1 and 5, "
27 "rounded to two decimal places, with 1 representing very poor quality and 5 representing excellent quality."
28 )
29
30 QUESTION_TEMPLATE = "{Question} Please only output the final answer with only one score in <answer> </answer> tags."
31
32 messages = []
33 for img_path in image_paths:
34 message = [
35 {
36 "role": "user",
37 "content": [
38 {'type': 'image', 'image': img_path},
39 {"type": "text", "text": QUESTION_TEMPLATE.format(Question=PROMPT)}
40 ],
41 }
42 ]
43 messages.append(message)
44
45 BSZ = 32
46 all_outputs = [] # List to store all answers
47 for i in tqdm(range(0, len(messages), BSZ)):
48 batch_messages = messages[i:i + BSZ]
49
50 # Preparation for inference
51 text = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True, add_vision_id=True) for msg in batch_messages]
52
53 image_inputs, video_inputs = process_vision_info(batch_messages)
54 inputs = processor(
55 text=text,
56 images=image_inputs,
57 videos=video_inputs,
58 padding=True,
59 return_tensors="pt",
60 )
61 inputs = inputs.to(device)
62
63 # Inference: Generation of the output
64 generated_ids = model.generate(**inputs, use_cache=True, max_new_tokens=512, do_sample=True, top_k=50, top_p=1)
65 generated_ids_trimmed = [
66 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
67 ]
68 batch_output_text = processor.batch_decode(
69 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
70 )
71
72 all_outputs.extend(batch_output_text)
73
74 path_score_dict = {}
75 for img_path, model_output in zip(image_paths, all_outputs):
76 try:
77 model_output_matches = re.findall(r'<answer>(.*?)</answer>', model_output, re.DOTALL)
78 model_answer = model_output_matches[-1].strip() if model_output_matches else model_output.strip()
79 score = float(re.search(r'\d+(\.\d+)?', model_answer).group())
80 except:
81 print(f"Meet error with {img_path}, please generate again.")
82 score = random.randint(1, 5)
83
84 path_score_dict[img_path] = score
85
86 return path_score_dict
87
88
89random.seed(1)
90MODEL_PATH = ""
91device = torch.device("cuda:3") if torch.cuda.is_available() else torch.device("cpu")
92
93model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
94 MODEL_PATH,
95 torch_dtype=torch.bfloat16,
96 attn_implementation="flash_attention_2",
97 device_map=device,
98)
99processor = AutoProcessor.from_pretrained(MODEL_PATH)
100processor.tokenizer.padding_side = "left"
101
102image_root = ""
103image_paths = get_image_paths(image_root) # It should be a list
104
105path_score_dict = score_batch_image(
106 image_paths, model, processor
107)
108
109file_name = "output.txt"
110with open(file_name, "w") as file:
111 for key, value in path_score_dict.items():
112 file.write(f"{key} {value}\n")
113
114print("Done!")1from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
2from qwen_vl_utils import process_vision_info
3
4import torch
5import random
6import re
7import os
8
9
10def score_image(image_path, model, processor):
11 PROMPT = (
12 "You are doing the image quality assessment task. Here is the question: "
13 "What is your overall rating on the quality of this picture? The rating should be a float between 1 and 5, "
14 "rounded to two decimal places, with 1 representing very poor quality and 5 representing excellent quality."
15 )
16
17 QUESTION_TEMPLATE = "{Question} First output the thinking process in <think> </think> tags and then output the final answer with only one score in <answer> </answer> tags."
18 # QUESTION_TEMPLATE = "Please describe the quality of this image."
19 message = [
20 {
21 "role": "user",
22 "content": [
23 {'type': 'image', 'image': image_path},
24 {"type": "text", "text": QUESTION_TEMPLATE.format(Question=PROMPT)}
25 ],
26 }
27 ]
28
29 batch_messages = [message]
30
31 # Preparation for inference
32 text = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True, add_vision_id=True) for msg in batch_messages]
33 image_inputs, video_inputs = process_vision_info(batch_messages)
34 inputs = processor(
35 text=text,
36 images=image_inputs,
37 videos=video_inputs,
38 padding=True,
39 return_tensors="pt",
40 )
41 inputs = inputs.to(device)
42
43 # Inference: Generation of the output
44 generated_ids = model.generate(**inputs, use_cache=True, max_new_tokens=2048, do_sample=True, top_k=50, top_p=1)
45 generated_ids_trimmed = [
46 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
47 ]
48 batch_output_text = processor.batch_decode(
49 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
50 )
51
52 reasoning = re.findall(r'<think>(.*?)</think>', batch_output_text[0], re.DOTALL)
53 reasoning = reasoning[-1].strip()
54
55 try:
56 model_output_matches = re.findall(r'<answer>(.*?)</answer>', batch_output_text[0], re.DOTALL)
57 model_answer = model_output_matches[-1].strip() if model_output_matches else batch_output_text[0].strip()
58 score = float(re.search(r'\d+(\.\d+)?', model_answer).group())
59 except:
60 print(f"================= Meet error with {img_path}, please generate again. =================")
61 score = random.randint(1, 5)
62
63 return reasoning, score
64
65
66random.seed(1)
67MODEL_PATH = ""
68device = torch.device("cuda:5") if torch.cuda.is_available() else torch.device("cpu")
69image_path = ""
70
71model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
72 MODEL_PATH,
73 torch_dtype=torch.bfloat16,
74 attn_implementation="flash_attention_2",
75 device_map=device,
76)
77processor = AutoProcessor.from_pretrained(MODEL_PATH)
78processor.tokenizer.padding_side = "left"
79
80reasoning, score = score_image(
81 image_path, model, processor
82)
83
84print(reasoning)
85print(score)1from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
2from qwen_vl_utils import process_vision_info
3from tqdm import tqdm
4
5import torch
6import random
7import re
8import os
9
10
11def get_image_paths(folder_path):
12 image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}
13 image_paths = []
14
15 for root, dirs, files in os.walk(folder_path):
16 for file in files:
17 _, ext = os.path.splitext(file)
18 if ext.lower() in image_extensions:
19 image_paths.append(os.path.join(root, file))
20
21 return image_paths
22
23def score_batch_image(image_paths, model, processor):
24 PROMPT = (
25 "You are doing the image quality assessment task. Here is the question: "
26 "What is your overall rating on the quality of this picture? The rating should be a float between 1 and 5, "
27 "rounded to two decimal places, with 1 representing very poor quality and 5 representing excellent quality."
28 )
29
30 QUESTION_TEMPLATE = "{Question} First output the thinking process in <think> </think> tags and then output the final answer with only one score in <answer> </answer> tags."
31
32 messages = []
33 for img_path in image_paths:
34 message = [
35 {
36 "role": "user",
37 "content": [
38 {'type': 'image', 'image': img_path},
39 {"type": "text", "text": QUESTION_TEMPLATE.format(Question=PROMPT)}
40 ],
41 }
42 ]
43 messages.append(message)
44
45 BSZ = 32
46 all_outputs = [] # List to store all answers
47 for i in tqdm(range(0, len(messages), BSZ)):
48 batch_messages = messages[i:i + BSZ]
49
50 # Preparation for inference
51 text = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True, add_vision_id=True) for msg in batch_messages]
52
53 image_inputs, video_inputs = process_vision_info(batch_messages)
54 inputs = processor(
55 text=text,
56 images=image_inputs,
57 videos=video_inputs,
58 padding=True,
59 return_tensors="pt",
60 )
61 inputs = inputs.to(device)
62
63 # Inference: Generation of the output
64 generated_ids = model.generate(**inputs, use_cache=True, max_new_tokens=512, do_sample=True, top_k=50, top_p=1)
65 generated_ids_trimmed = [
66 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
67 ]
68 batch_output_text = processor.batch_decode(
69 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
70 )
71
72 all_outputs.extend(batch_output_text)
73
74 path_score_dict = {}
75 for img_path, model_output in zip(image_paths, all_outputs):
76 reasoning = re.findall(r'<think>(.*?)</think>', model_output, re.DOTALL)
77 reasoning = reasoning[-1].strip()
78
79 try:
80 model_output_matches = re.findall(r'<answer>(.*?)</answer>', model_output, re.DOTALL)
81 model_answer = model_output_matches[-1].strip() if model_output_matches else model_output.strip()
82 score = float(re.search(r'\d+(\.\d+)?', model_answer).group())
83 except:
84 print(f"Meet error with {img_path}, please generate again.")
85 score = random.randint(1, 5)
86
87 path_score_dict[img_path] = score
88
89 return path_score_dict
90
91
92random.seed(1)
93MODEL_PATH = ""
94device = torch.device("cuda:3") if torch.cuda.is_available() else torch.device("cpu")
95
96model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
97 MODEL_PATH,
98 torch_dtype=torch.bfloat16,
99 attn_implementation="flash_attention_2",
100 device_map=device,
101)
102processor = AutoProcessor.from_pretrained(MODEL_PATH)
103processor.tokenizer.padding_side = "left"
104
105image_root = ""
106image_paths = get_image_paths(image_root) # It should be a list
107
108path_score_dict = score_batch_image(
109 image_paths, model, processor
110)
111
112file_name = "output.txt"
113with open(file_name, "w") as file:
114 for key, value in path_score_dict.items():
115 file.write(f"{key} {value}\n")
116
117print("Done!")1# Please install vLLM first: https://docs.vllm.ai/en/stable/getting_started/installation/gpu.html
2
3from transformers import Qwen2_5_VLProcessor, AutoProcessor
4from vllm import LLM, RequestOutput, SamplingParams
5from qwen_vl_utils import process_vision_info
6
7import torch
8import random
9import re
10import os
11
12IMAGE_PATH = "./images"
13MODEL_PATH = "TianheWu/VisualQuality-R1-7B"
14
15def get_image_paths(folder_path):
16 image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}
17 image_paths = []
18
19 for root, dirs, files in os.walk(folder_path):
20 for file in files:
21 _, ext = os.path.splitext(file)
22 if ext.lower() in image_extensions:
23 image_paths.append(os.path.join(root, file))
24
25 return image_paths
26
27def score_batch_image(image_paths, model: LLM, processor: Qwen2_5_VLProcessor):
28 PROMPT = (
29 "You are doing the image quality assessment task. Here is the question: "
30 "What is your overall rating on the quality of this picture? The rating should be a float between 1 and 5, "
31 "rounded to two decimal places, with 1 representing very poor quality and 5 representing excellent quality."
32 )
33
34 QUESTION_TEMPLATE = "{Question} First output the thinking process in <think> </think> tags and then output the final answer with only one score in <answer> </answer> tags."
35
36 messages = []
37 for img_path in image_paths:
38 message = [
39 {
40 "role": "user",
41 "content": [
42 {'type': 'image', 'image': img_path},
43 {"type": "text", "text": QUESTION_TEMPLATE.format(Question=PROMPT)}
44 ],
45 }
46 ]
47 messages.append(message)
48
49 all_outputs = [] # List to store all answers
50
51 # Preparation for inference
52 print("preprocessing ...")
53 texts = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True, add_vision_id=True) for msg in messages]
54 image_inputs, video_inputs = process_vision_info(messages)
55
56 inputs = [{
57 "prompt": texts[i],
58 "multi_modal_data": {
59 "image": image_inputs[i]
60 },
61 } for i in range(len(messages))]
62
63 output: list[RequestOutput] = model.generate(
64 inputs,
65 sampling_params=SamplingParams(
66 max_tokens=512,
67 temperature=0.1,
68 top_k=50,
69 top_p=1.0,
70 stop_token_ids=[processor.tokenizer.eos_token_id],
71 ),
72 )
73
74 batch_output_text = [o.outputs[0].text for o in output]
75
76 all_outputs.extend(batch_output_text)
77
78 path_score_dict = {}
79 for img_path, model_output in zip(image_paths, all_outputs):
80 print(f"{model_output = }")
81 try:
82 model_output_matches = re.findall(r'<answer>(.*?)</answer>', model_output, re.DOTALL)
83 model_answer = model_output_matches[-1].strip() if model_output_matches else model_output.strip()
84 score = float(re.search(r'\d+(\.\d+)?', model_answer).group())
85 except:
86 print(f"Meet error with {img_path}, please generate again.")
87 score = random.randint(1, 5)
88
89 path_score_dict[img_path] = score
90
91 return path_score_dict
92
93
94random.seed(1)
95model = LLM(
96 model=MODEL_PATH,
97 tensor_parallel_size=1,
98 trust_remote_code=True,
99 seed=1,
100)
101
102processor = AutoProcessor.from_pretrained(MODEL_PATH)
103processor.tokenizer.padding_side = "left"
104
105image_paths = get_image_paths(IMAGE_PATH) # It should be a list
106
107path_score_dict = score_batch_image(
108 image_paths, model, processor
109)
110
111file_name = "output.txt"
112with open(file_name, "w") as file:
113 for key, value in path_score_dict.items():
114 file.write(f"{key} {value}\n")
115
116print("Done!")cd datasets, then run python make_data.py (with moderate modifications) to generate a JSON file for model training.src/open-r1-multimodal/run_scripts/KADID-10K/one_node_run_kadid.sh:--model_name_or_path [Your Qwen2.5-VL-7B-Instruct path] \
--image_folders [Your dataset images path] \
--data_file_paths [Your JSON file path] \bash src/open-r1-multimodal/run_scripts/KADID-10K/one_node_run_kadid.shbash src/open-r1-multimodal/run_scripts/KADID-10K/multi_run_kadid.shsigstianhewu@gmail.com or tianhewu-c@my.cityu.edu.hk.@article{wu2025visualquality,
title={{VisualQuality-R1}: Reasoning-Induced Image Quality Assessment via Reinforcement Learning to Rank},
author={Wu, Tianhe and Zou, Jian and Liang, Jie and Zhang, Lei and Ma, Kede},
journal={arXiv preprint arXiv:2505.14460},
year={2025}
}