Upload 5 files

README.md

@@ -1,14 +1,136 @@
----
-title: FGResQ
-emoji: 🐢
-colorFrom: purple
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: FGResQ office demo
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<div align="center">
+    <a href="https://arxiv.org/abs/2508.14475"><img src="https://img.shields.io/badge/Arxiv-preprint-red"></a>
+    <a href="https://pxf0429.github.io/FGResQ/"><img src="https://img.shields.io/badge/Homepage-green"></a>
+    <a href='https://github.com/sxfly99/FGRestore/stargazers'><img src='https://img.shields.io/github/stars/sxfly99/FGRestore.svg?style=social'></a>
+</div>
+
+<h1 align="center">Fine-grained Image Quality Assessment for Perceptual Image Restoration</h1>
+
+<div align="center">
+    <a href="https://github.com/sxfly99">Xiangfei Sheng</a><sup>1*</sup>,
+    <a href="https://github.com/pxf0429">Xiaofeng Pan</a><sup>1*</sup>,
+    <a href="https://github.com/yzc-ippl">Zhichao Yang</a><sup>1</sup>,
+    <a href="https://faculty.xidian.edu.cn/cpf/">Pengfei Chen</a><sup>1</sup>,
+    <a href="https://web.xidian.edu.cn/ldli/">Leida Li</a><sup>1#</sup>
+</div>
+
+<div align="center">
+  <sup>1</sup>School of Artificial Intelligence, Xidian University
+</div>
+
+<div align="center">
+<sup>*</sup>Equal contribution. <sup>#</sup>Corresponding author. 
+</div>
+
+
+<div align="center">
+  <img src="FGResQ.png" width="800"/>
+</div>
+
+<div style="font-family: sans-serif; margin-bottom: 2em;">
+    <h2 style="border-bottom: 1px solid #eaecef; padding-bottom: 0.3em; margin-bottom: 1em;">📰 News</h2>
+    <ul style="list-style-type: none; padding-left: 0;">
+        <li style="margin-bottom: 0.8em;">
+            <strong>[2025-11-08]</strong> 🎉🎉🎉Our paper, "Fine-grained Image Quality Assessment for Perceptual Image Restoration", has been accepted to appear at AAAI 2026!
+        </li>
+        <li style="margin-bottom: 0.8em;">
+            <strong>[2025-08-20]</strong> Code and pre-trained models for FGResQ released.
+        </li>
+    </ul>
+</div>
+
+
+## Quick Start
+
+This guide will help you get started with the FGResQ inference code.
+
+### 1. Installation
+
+First, clone the repository and install the required dependencies.
+
+```bash
+git clone https://github.com/sxfly99/FGResQ.git
+cd FGResQ
+pip install -r requirements.txt
+```
+
+### 2. Download Pre-trained Weights
+
+You can download the pre-trained model weights from the following link:
+[**Download Weights (Google Drive)**](https://drive.google.com/drive/folders/10MVnAoEIDZ08Rek4qkStGDY0qLiWUahJ?usp=drive_link) or [**(Baidu Netdisk)**](https://pan.baidu.com/s/1a2IZbr_PrgZYCbUbjKLykA?pwd=9ivu)
+
+Place the downloaded files in the `weights` directory.
+
+- `FGResQ.pth`: The main model for quality scoring and ranking.
+- `Degradation.pth`: The weights for the degradation-aware task branch.
+
+Create the `weights` directory if it doesn't exist and place the files inside.
+
+```
+FGRestore/
+|-- weights/
+|   |-- FGResQ.pth
+|   |-- Degradation.pth
+|-- model/
+|   |-- FGResQ.py
+|-- requirements.txt
+|-- README.md
+```
+
+## Usage
+
+The `FGResQ` provides two main functionalities: scoring a single image and comparing a pair of images.
+
+### Initialize the Scorer
+
+First, import and initialize the `FGResQ`.
+
+```python
+from model.FGResQ import FGResQ
+
+# Path to the main model weights
+model_path = "weights/FGResQ.pth"
+
+# Initialize the inference engine
+model = FGResQ(model_path=model_path)
+```
+
+### 1. Single Image Input Mode: Quality Scoring
+
+You can get a quality score for a single image. The score typically ranges from 0 to 1, where a higher score indicates better quality.
+
+```python
+image_path = "path/to/your/image.jpg"
+quality_score = model.predict_single(image_path)
+print(f"The quality score for the image is: {quality_score:.4f}")
+```
+
+### 2. Pairwise Image Input Mode: Quality Ranking
+
+You can also compare two images to determine which one has better quality.
+
+```python
+image_path1 = "path/to/image1.jpg"
+image_path2 = "path/to/image2.jpg"
+
+comparison_result = model.predict_pair(image_path1, image_path2)
+
+# The result includes a human-readable comparison and raw probabilities
+print(f"Comparison: {comparison_result['comparison']}")
+# Example output: "Comparison: Image 1 is better"
+
+print(f"Raw output probabilities: {comparison_result['comparison_raw']}")
+# Example output: "[0.8, 0.1, 0.1]" (Probabilities for Image1 > Image2, Image2 > Image1, Image1 ≈ Image2)
+```
+## Citation
+
+If you find this work is useful, pleaes cite our paper!
+
+```bibtex
+
+@article{sheng2025fgresq,
+  title={Fine-grained Image Quality Assessment for Perceptual Image Restoration},
+  author={Sheng, Xiangfei and Pan, Xiaofeng and Yang, Zhichao and Chen, Pengfei and Li, Leida},
+  journal={arXiv preprint arXiv:2508.14475},
+  year={2025}
+}

model/FGResQ.py

@@ -0,0 +1,271 @@
+import numpy as np
+import timm
+import torch
+import torchvision
+import torch.nn as nn
+from transformers import CLIPVisionModel
+# import open_clip
+import torchvision.transforms as transforms
+from PIL import Image
+import cv2
+import accelerate
+
+def load_clip_model(clip_model="openai/ViT-B-16", clip_freeze=True, precision='fp16'):
+    pretrained, model_tag = clip_model.split('/')
+    pretrained = None if pretrained == 'None' else pretrained
+    # clip_model = open_clip.create_model(model_tag, precision=precision, pretrained=pretrained)
+    # clip_model = timm.create_model('timm/vit_base_patch16_clip_224.openai', pretrained=True, in_chans=3)
+    clip_model = CLIPVisionModel.from_pretrained(clip_model)
+    if clip_freeze:
+        for param in clip_model.parameters():
+            param.requires_grad = False
+
+    if model_tag == 'clip-vit-base-patch16':
+        feature_size = dict(global_feature=768, local_feature=[196, 768])
+    elif model_tag == 'ViT-L-14-quickgelu' or model_tag == 'ViT-L-14':
+        feature_size = dict(global_feature=768, local_feature=[256, 1024])
+    else:
+        raise ValueError(f"Unknown model_tag: {model_tag}")
+
+    return clip_model, feature_size
+
+class DualBranch(nn.Module):
+
+    def __init__(self, clip_model="openai/clip-vit-base-patch16", clip_freeze=True, precision='fp16'):
+        super(DualBranch, self).__init__()
+        self.clip_freeze = clip_freeze
+
+        # Load CLIP model
+        self.clip_model, feature_size = load_clip_model(clip_model, clip_freeze, precision)
+
+        # Initialize CLIP vision model for task classification
+        self.task_cls_clip = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch16")
+        
+
+        self.head = nn.Linear(feature_size['global_feature']*3, 1)
+        self.compare_head =nn.Linear(feature_size['global_feature']*6, 3)
+        
+    
+        self.prompt = nn.Parameter(torch.rand(1, feature_size['global_feature']))
+        self.task_mlp = nn.Sequential(
+            nn.Linear(feature_size['global_feature'], feature_size['global_feature']), 
+            nn.SiLU(False),
+            nn.Linear(feature_size['global_feature'], feature_size['global_feature']))
+        self.prompt_mlp = nn.Linear(feature_size['global_feature'], feature_size['global_feature'])
+        
+        with torch.no_grad():
+            self.task_mlp[0].weight.fill_(0.0)
+            self.task_mlp[0].bias.fill_(0.0)
+            self.task_mlp[2].weight.fill_(0.0)
+            self.task_mlp[2].bias.fill_(0.0)
+            self.prompt_mlp.weight.fill_(0.0)
+            self.prompt_mlp.bias.fill_(0.0)
+        
+        # Load pre-trained weights
+        self._load_pretrained_weights("./weights/Degradation.pth")
+
+
+        for param in self.task_cls_clip.parameters():
+            param.requires_grad = False
+
+        # Unfreeze the last two layers
+        for i in range(10, 12):  # Layers 10 and 11
+            for param in self.task_cls_clip.vision_model.encoder.layers[i].parameters():
+                param.requires_grad = True
+    def _load_pretrained_weights(self, state_dict_path):
+        """
+        Load pre-trained weights, including the CLIP model and classification head.
+        """
+        # Load state dictionary
+        state_dict = torch.load(state_dict_path)
+        
+        # Separate weights for CLIP model and classification head
+        clip_state_dict = {}
+
+        for key, value in state_dict.items():
+            if key.startswith('clip_model.'):
+                # Remove 'clip_model.' prefix for the CLIP model
+                new_key = key.replace('clip_model.', '')
+                clip_state_dict[new_key] = value
+            # elif key in ['head.weight', 'head.bias']:
+            #     # Save weights for the classification head
+            #     head_state_dict[key] = value
+        
+        # Load weights for the CLIP model
+        self.task_cls_clip.load_state_dict(clip_state_dict, strict=False)
+        print("Successfully loaded CLIP model weights")
+        
+    def forward(self, x0, x1 = None):
+        # features, _ = self.clip_model.encode_image(x)
+        if x1 is None:
+            # Image features
+            features0 = self.clip_model(x0)['pooler_output']
+            # Classification features
+            task_features0 = self.task_cls_clip(x0)['pooler_output']
+
+            # Learn classification features
+            task_embedding = torch.softmax(self.task_mlp(task_features0), dim=1) * self.prompt
+            task_embedding = self.prompt_mlp(task_embedding)
+
+            # features = torch.cat([features0, task_features], dim
+            features0 = torch.cat([features0, task_embedding, features0+task_embedding], dim=1)
+            quality = self.head(features0)
+            quality = nn.Sigmoid()(quality)
+
+            return quality, None, None
+        elif x1 is not None:
+            # features_, _ = self.clip_model.encode_image(x_local)
+            # Image features
+            features0 = self.clip_model(x0)['pooler_output']
+            features1 = self.clip_model(x1)['pooler_output']
+            # Classification features
+            task_features0 = self.task_cls_clip(x0)['pooler_output']
+            task_features1 = self.task_cls_clip(x1)['pooler_output']
+
+            task_embedding0 = torch.softmax(self.task_mlp(task_features0), dim=1) * self.prompt
+            task_embedding0 = self.prompt_mlp(task_embedding0)
+            task_embedding1 = torch.softmax(self.task_mlp(task_features1), dim=1) * self.prompt
+            task_embedding1 = self.prompt_mlp(task_embedding1)
+
+            features0 = torch.cat([features0, task_embedding0, features0+task_embedding0], dim=1)
+            features1 = torch.cat([features1, task_embedding1, features1+task_embedding1], dim=1)
+
+            # features0 = torch.cat([features0, task_features0], dim=
+            # import pdb; pdb.set_trace()
+            features = torch.cat([features0, features1], dim=1)
+            # features = torch.cat([features0, features1], dim=1)
+            compare_quality = self.compare_head(features)
+
+            # quality0 = self.head(features0)
+            # quality1 = self.head(features1)
+            quality0 = self.head(features0)
+            quality1 = self.head(features1)
+            quality0 = nn.Sigmoid()(quality0)
+            quality1 = nn.Sigmoid()(quality1)
+
+            # quality = {'quality0': quality0, 'quality1': quality1}
+
+            return quality0, quality1, compare_quality
+
+class FGResQ:
+    def __init__(self, model_path, clip_model="openai/clip-vit-base-patch16", input_size=224, device=None):
+        """
+        Initializes the inference model.
+
+        Args:
+            model_path (str): Path to the pre-trained model checkpoint (.pth or .safetensors).
+            clip_model (str): Name of the CLIP model to use.
+            input_size (int): Input image size for the model.
+            device (str, optional): Device to run inference on ('cuda' or 'cpu'). Auto-detected if None.
+        """
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+        
+        print(f"Using device: {self.device}")
+
+        # Load the model
+        self.model = DualBranch(clip_model=clip_model, clip_freeze=True, precision='fp32')
+        # self.model = self.accelerator.unwrap_model(self.model)
+        # Load model weights
+        try:
+            raw = torch.load(model_path, map_location=self.device)
+            # unwrap possible containers
+            if isinstance(raw, dict) and any(k in raw for k in ['model', 'state_dict']):
+                state_dict = raw.get('model', raw.get('state_dict', raw))
+            else:
+                state_dict = raw
+
+            # Only strip 'module.' if present; keep other namespaces intact
+            if any(k.startswith('module.') for k in state_dict.keys()):
+                state_dict = {k.replace('module.', '', 1): v for k, v in state_dict.items()}
+
+            missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
+            if missing:
+                print(f"[load_state_dict] Missing keys: {missing}")
+            if unexpected:
+                print(f"[load_state_dict] Unexpected keys: {unexpected}")
+            print(f"Model weights loaded from {model_path}")
+        except Exception as e:
+            print(f"Error loading model weights: {e}")
+            raise
+
+        self.model.to(self.device)
+        self.model.eval()
+
+        # Define image preprocessing
+        # Match training/validation pipeline: first unify to 256x256 (as in cls_model/dataset.py),
+        # then CenterCrop to input_size, followed by CLIP normalization.
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.CenterCrop(input_size),
+            transforms.Normalize(
+                mean=[0.48145466, 0.4578275, 0.40821073],
+                std=[0.26862954, 0.26130258, 0.27577711]
+            )
+        ])
+
+    def _preprocess_image(self, image_path):
+        """Load and preprocess a single image."""
+        try:
+            # Match training dataset loader: cv2 read + resize to 256x256 (INTER_LINEAR)
+            img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+            if img is None:
+                raise FileNotFoundError(f"Failed to read image at {image_path}")
+            img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_LINEAR)
+            image = Image.fromarray(img)
+            image_tensor = self.transform(image).unsqueeze(0)
+            return image_tensor.to(self.device)
+        except FileNotFoundError:
+            print(f"Error: Image file not found at {image_path}")
+            return None
+        except Exception as e:
+            print(f"Error processing image {image_path}: {e}")
+            return None
+
+    @torch.no_grad()
+    def predict_single(self, image_path):
+        """
+        Predict the quality score of a single image.
+        """
+        image_tensor = self._preprocess_image(image_path)
+        if image_tensor is None:
+            return None
+
+        quality_score, _, _ = self.model(image_tensor)
+        return quality_score.squeeze().item()
+
+    @torch.no_grad()
+    def predict_pair(self, image_path1, image_path2):
+        """
+        Compare the quality of two images.
+        """
+        image_tensor1 = self._preprocess_image(image_path1)
+        image_tensor2 = self._preprocess_image(image_path2)
+
+        if image_tensor1 is None or image_tensor2 is None:
+            return None
+
+        quality1, quality2, compare_result = self.model(image_tensor1, image_tensor2)
+        
+        quality1 = quality1.squeeze().item()
+        quality2 = quality2.squeeze().item()
+        
+        # Interpret the comparison result
+        # print(compare_result.shape)
+        compare_probs = torch.softmax(compare_result, dim=-1).squeeze(dim=0).cpu().numpy()
+        # print(compare_probs)
+        prediction = np.argmax(compare_probs)
+
+        # Align with training label semantics:
+        # dataset encodes prefs: A>B -> 1, A<B -> 0, equal -> 2
+        # So class 1 => Image 1 (A) is better, class 0 => Image 2 (B) is better
+        comparison_map = {0: 'Image 2 is better', 1: 'Image 1 is better', 2: 'Images are of similar quality'}
+
+        return {
+            'comparison': comparison_map[prediction],
+            'comparison_raw': compare_probs.tolist()}
+    
+    
+        

requirements.txt

@@ -0,0 +1,46 @@
+accelerate==0.28.0
+certifi==2025.7.14
+charset-normalizer==3.4.2
+contourpy==1.1.1
+cycler==0.12.1
+einops==0.8.1
+filelock==3.16.1
+fonttools==4.57.0
+fsspec==2025.3.0
+hf-xet==1.1.5
+huggingface-hub==0.33.4
+idna==3.10
+importlib_resources==6.4.5
+kiwisolver==1.4.7
+matplotlib==3.7.5
+mpmath==1.3.0
+numpy==1.24.4
+opencv-python==4.8.1.78
+packaging==25.0
+pandas==2.0.3
+pillow==10.4.0
+prefetch-generator==1.0.3
+psutil==7.0.0
+pyparsing==3.1.4
+python-dateutil==2.9.0.post0
+pytz==2025.2
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.4
+safetensors==0.5.3
+scipy==1.10.1
+seaborn==0.13.2
+six==1.17.0
+sympy==1.12
+timm==1.0.17
+tokenizers==0.15.2
+torch==1.13.0+cu117
+torch-ema==0.3
+torchaudio==0.13.0+cu117
+torchvision==0.14.0+cu117
+tqdm==4.67.1
+transformers==4.36.1
+typing_extensions==4.13.2
+tzdata==2025.2
+urllib3==2.2.3
+zipp==3.20.2