test

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+JarvisIR/package/agent_tools/RIDCP/.eggs/**/*.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+**/.eggs/
+*.so
+JarvisIR/package/agent_tools/RIDCP/.eggs

JarvisIR/.gitignore

@@ -0,0 +1,146 @@
+# ignored folders
+datasets/*
+experiments*
+experiments/*
+# results/*
+tb_logger*
+pretrained_models
+wandb/*
+tmp*/*
+*.sh
+.vscode*
+.github
+# ignored files
+version.py
+
+# ignored files with suffix
+*.html
+*.pth
+*.zip
+
+# template
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+checkpoint
+*.zip
+*.png
+*.jpg
+*.jpeg
+*.json
+.eggs
+*.pth*
+*.pk
+*.pkl
+*.mdb
+*.pt
+*.log
+*.bin
+__pycache__
+Temp
+work_dirs
+src/sft/config
+checkpoints

JarvisIR/.gitmodules

@@ -0,0 +1,7 @@
+[submodule "dependences/BasicSR"]
+	path = dependences/BasicSR
+	url = https://github.com/XPixelGroup/BasicSR.git
+[submodule "src/sft/xtuner"]
+	path = src/sft/xtuner
+	url = https://github.com/InternLM/xtuner.git
+	branch = v0.1.23

JarvisIR/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Yunlong Lin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

JarvisIR/README.md

@@ -0,0 +1,148 @@
+<div align="center">
+<div align="center">
+  <img src="assets/icon.png" alt="JarvisIR Logo" width="100px">
+  <h1>[CVPR' 2025] JarvisIR: Elevating Autonomous Driving Perception with Intelligent Image Restoration</h1>
+</div>
+
+<a href="https://lyl1015.github.io/papers/CVPR2025_JarvisIR.pdf" target="_blank" rel="noopener noreferrer">
+  <img src="https://img.shields.io/badge/Paper-JarvisIR-b31b1b" alt="Paper PDF">
+</a>
+<!-- <a href="#"><img src="https://img.shields.io/badge/arXiv-即将发布-b31b1b" alt="arXiv"></a> -->
+<a href="https://cvpr2025-jarvisir.github.io/"><img src="https://img.shields.io/badge/Project-Page-green" alt="Project Page"></a>
+<a href="#"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Coming%20Soon-blue" alt="Demo"></a>
+<a href="https://github.com/LYL1015/JarvisIR?tab=readme-ov-file/"><img src="https://img.shields.io/badge/GitHub-Code-black" alt="Code"></a>
+
+
+[Yunlong Lin](https://lyl1015.github.io/)<sup>1*♣</sup>, [Zixu Lin](https://github.com/)<sup>1*♣</sup>, [Haoyu Chen](https://haoyuchen.com/)<sup>2*</sup>, [Panwang Pan](https://paulpanwang.github.io/)<sup>3*</sup>, [Chenxin Li](https://chenxinli001.github.io/)<sup>6</sup>, [Sixiang Chen](https://ephemeral182.github.io/)<sup>2</sup>, [Kairun Wen](https://kairunwen.github.io/)<sup>1</sup>, [Yeying Jin](https://jinyeying.github.io/)<sup>4</sup>, [Wenbo Li](https://fenglinglwb.github.io/)<sup>5†</sup>, [Xinghao Ding](https://scholar.google.com/citations?user=k5hVBfMAAAAJ&hl=zh-CN)<sup>1†</sup>
+
+<sup>1</sup>Xiamen University, <sup>2</sup>The Hong Kong University of Science and Technology (Guangzhou), <sup>3</sup>Bytedance's Pico, <sup>4</sup>Tencent, <sup>5</sup>Huawei Noah's Ark Lab, <sup>6</sup>The Chinese University of Hong Kong
+<!-- <sup>*</sup>Equal Contribution <sup>♣</sup>Equal Contribution <sup>†</sup>Corresponding Author -->
+Accepted by CVPR 2025
+
+<!-- <div align="center">
+  <video width="800" controls>
+    <source src="assets/demo.mp4" type="video/mp4">
+    Your browser does not support the video tag.
+  </video>
+  <p>JarvisIR Demo Video: Showcasing image restoration capabilities under various adverse weather conditions</p>
+</div> -->
+https://github.com/user-attachments/assets/d9094fba-e24c-403e-90cb-b3d2f6e48939
+
+
+</div>
+
+
+## :postbox: Updates
+<!-- - 2023.12.04: Add an option to speed up the inference process by adjusting the number of denoising steps. -->
+<!-- - 2024.2.9: Release our demo codes and models. Have fun! :yum: -->
+- 2025.4.8: This repo is created.
+
+## :diamonds: Overview
+JarvisIR (CVPR 2025) is a VLM-powered agent designed to tackle the challenges of vision-centric perception systems under unpredictable and coupled weather degradations. It leverages the VLM as a controller to manage multiple expert restoration models, enabling robust and autonomous operation in real-world conditions. JarvisIR employs a novel two-stage framework consisting of supervised fine-tuning and human feedback alignment, allowing it to effectively fine-tune on large-scale real-world data in an unsupervised manner. Supported by CleanBench, a comprehensive dataset with 150K synthetic and 80K real instruction-response pairs, JarvisIR demonstrates superior decision-making and restoration capabilities, achieving a 50% improvement in the average of all perception metrics on CleanBench-Real.
+<div align="center">
+  <img src="assets/teaser1.png" alt="JarvisIR Teaser" width="800px">
+</div>
+
+## :rocket: Method
+
+JarvisIR implements an innovative two-stage framework that leverages a Vision-Language Model (VLM) as a controller to manage multiple expert restoration models:
+
+1. **Supervised Fine-tuning Stage**: JarvisIR undergoes supervised fine-tuning on synthetic data from CleanBench to enable it to follow user instructions and recognize image degradation. This initial training allows the model to identify various types of image degradation and select appropriate restoration strategies.
+
+2. **Human Feedback Alignment Stage**: We further finetune JarvisIR on CleanBench-Real using the MRRHF algorithm to improve system robustness, reduce hallucinations, and enhance generalizability under real-world adverse weather conditions. This stage ensures the model makes decisions that align with human expectations in complex real-world scenarios.
+
+The core advantage of JarvisIR lies in its ability to handle multiple complex, coupled weather degradations and provide stable, reliable image inputs for autonomous driving perception systems.
+
+<div align="center">
+  <img src="assets/framework.png" alt="JarvisIR Method" width="800px">
+  <p>Two-stage training framework of JarvisIR: The VLM controller analyzes input images, selects and coordinates expert models for optimal restoration</p>
+</div>
+
+## :bar_chart: CleanBench Dataset
+
+To support the training and evaluation of JarvisIR, we introduce CleanBench, the first high-quality instruction-following dataset specifically curated for developing intelligent restoration systems. CleanBench contains **150K** synthetic and **80K** real instruction-response pairs, providing a comprehensive foundation for training and evaluating intelligent image restoration systems.
+
+### Dataset Construction
+
+The CleanBench dataset construction workflow consists of three main steps:
+
+1. **Synthesis of Degraded Images**: We generate a diverse set of degraded images by applying various weather conditions and degradation types to clean images, creating realistic scenarios that autonomous driving systems might encounter.
+
+2. **Generation of Assessment Reasoning and Optimal Task Sequence**: For each degraded image, we generate detailed assessments of the degradation types present and determine the optimal sequence of restoration tasks needed to effectively restore the image.
+
+3. **Generation of Instruction-Response Pairs**: Based on the degradation assessment and restoration sequence, we create comprehensive instruction-response pairs that guide the model in understanding user requests and providing appropriate restoration solutions.
+
+<div align="center">
+  <img src="assets/datasets.png" alt="CleanBench Dataset Construction" width="800px">
+  <p>CleanBench dataset construction workflow: from degraded image synthesis to instruction-response pair generation</p>
+</div>
+
+### Dataset Features
+
+- **Comprehensive Coverage**: Includes various weather conditions (rain, snow, fog, night) and their combinations
+- **High-Quality Annotations**: Detailed degradation assessments and optimal restoration sequences
+- **Real-World Examples**: 80K real-world examples to ensure model generalization
+- **Instruction Diversity**: Multiple instruction formats to enhance model adaptability
+
+CleanBench serves as a crucial resource for training and evaluating intelligent image restoration systems, enabling models like JarvisIR to make informed decisions about restoration strategies in complex real-world scenarios.
+
+## :computer: Getting Started
+
+For sft training and environment setup preparation, please follow:
+
+- [SFT Training](./docs/sft_training.md)
+<!-- - [Dataset Preparation](./docs/dataset_preparation.md) -->
+<!-- 
+For sft training, please follow:
+
+- [SFT Training](./docs/sft_training.md) -->
+
+For gradio demo runing, please follow:
+
+- [Gradio Demo](./docs/gradio_demo.md)
+
+
+## :toolbox: Expert Models
+
+JarvisIR integrates multiple expert restoration models to handle various types of image degradation:
+
+| Task | Model | Description |
+|------|-------|-------------|
+| **Super-resolution** | Real-ESRGAN | Fast GAN-based model for super-resolution, deblurring, and artifact removal |
+| **Denoising** | SCUNet | Hybrid UNet-based model combining convolution and transformer blocks for robust denoising |
+| **Deraining** | UDR-S2Former | Uncertainty-aware transformer model for rain streak removal |
+|  | Img2img-turbo-rain | Efficient SD-turbo based model for fast and effective rain removal |
+| **Raindrop removal** | IDT | Transformer-based model for de-raining and raindrop removal |
+| **Dehazing** | RIDCP | Efficient dehazing model utilizing high-quality codebook priors |
+|  | KANet | Efficient dehazing network using a localization-and-removal pipeline |
+| **Desnowing** | Img2img-turbo-snow | Efficient model for removing snow artifacts while preserving natural scene details |
+|  | Snowmaster | Real-world image desnowing via MLLM with multi-model feedback optimization |
+| **Low-light enhancement** | Retinexformer | One-stage Retinex-based Transformer for low-light image enhancement |
+|  | HVICIDNet | Lightweight transformer for low-light and exposure correction |
+|  | LightenDiff | Diffusion-based framework for low-light enhancement |
+
+
+## :circus_tent: Checklist
+
+- [x] Release preview inference code and gradio demo
+- [x] Release SFT training code
+- [ ] Release Hugging Face demo
+- [ ] Release CleanBench data
+
+
+## :pray: Acknowledgements
+
+We would like to express our gratitude to [HuggingGPT](https://github.com/microsoft/JARVIS), [XTuner](https://github.com/InternLM/xtuner), and [RRHF](https://github.com/GanjinZero/RRHF) for their valuable open-source contributions which have provided important technical references for our work.
+
+
+
+
+## :love_you_gesture: Citation
+```bibtex
+@inproceedings{jarvisir2025,
+  title={JarvisIR: Elevating Autonomous Driving Perception with Intelligent Image Restoration},
+  author={Lin, Yunlong and Lin, Zixu and Chen, Haoyu and Pan, Panwang and Li, Chenxin and Chen, Sixiang and Kairun, Wen and Jin, Yeying and Li, Wenbo and Ding, Xinghao},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2025}
+}
+```

JarvisIR/demo_gradio.py

@@ -0,0 +1,474 @@
+import os
+import re
+import random
+import gradio as gr
+import torch
+from PIL import Image
+from transformers import AutoProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
+from threading import Thread
+from agent_tools import RestorationToolkit
+
+# Set CUDA device
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+# Model configuration
+# XXX: Path to the fine-tuned LLaVA model
+model_id = "" 
+    
+# Available image restoration tasks and their corresponding models
+all_tasks = " {denoise: [scunet, restormer], lighten: [retinexformer_fivek, hvicidnet, lightdiff], \
+                derain: [idt, turbo_rain, s2former], defog:[ridcp, kanet], \
+                desnow:[turbo_snow, snowmaster], super_resolution: [real_esrgan], \
+            }"
+
+# Various prompt templates for querying the LLM about image degradation and restoration tasks
+prompts_query2 = [
+    f"Considering the image's degradation, suggest the required tasks with explanations, and identify suitable tools for each task. Options for tasks and tools include: {all_tasks}.",
+    f"Given the image's degradation, outline the essential tasks along with justifications, and choose the appropriate tools for each task from the following options: {all_tasks}.",
+    f"Please specify the tasks required due to the image's degradation, explain the reasons, and select relevant tools for each task from the provided options: {all_tasks}.",
+    f"Based on the image degradation, determine the necessary tasks and their reasons, along with the appropriate tools for each task. Choose from these options: {all_tasks}.",
+    f"Identify the tasks required to address the image's degradation, including the reasons for each, and select tools from the options: {all_tasks}.",
+    f"Considering the degradation observed, list the tasks needed and their justifications, then pick the most suitable tools for each task from these options: {all_tasks}.",
+    f"Evaluate the image degradation, and based on that, provide the necessary tasks and reasons, along with tools chosen from the options: {all_tasks}.",
+    f"With respect to the image degradation, outline the tasks needed and explain why, selecting tools from the following list: {all_tasks}.",
+    f"Given the level of degradation in the image, specify tasks to address it, include reasons, and select tools for each task from: {all_tasks}.",
+    f"Examine the image's degradation, propose relevant tasks and their explanations, and identify tools from the options provided: {all_tasks}.",
+    f"Based on observed degradation, detail the tasks required, explain your choices, and select tools from these options: {all_tasks}.",
+    f"Using the image's degradation as a guide, list the necessary tasks, include explanations, and pick tools from the provided choices: {all_tasks}.",
+    f"Assess the image degradation, provide the essential tasks and reasons, and select the appropriate tools for each task from the options: {all_tasks}.",
+    f"According to the image's degradation, determine which tasks are necessary and why, choosing tools for each task from: {all_tasks}.",
+    f"Observe the degradation in the image, specify the needed tasks with justifications, and select appropriate tools from: {all_tasks}.",
+    f"Taking the image degradation into account, specify tasks needed, provide reasons, and choose tools from the following: {all_tasks}.",
+    f"Consider the image's degradation level, outline the tasks necessary, provide reasoning, and select suitable tools from: {all_tasks}.",
+    f"Evaluate the degradation in the image, identify tasks required, explain your choices, and pick tools from: {all_tasks}.",
+    f"Analyze the image degradation and suggest tasks with justifications, choosing the best tools from these options: {all_tasks}.",
+    f"Review the image degradation, and based on it, specify tasks needed, provide reasons, and select tools for each task from: {all_tasks}."
+]
+
+# Initialize models
+print("Loading LLM model...")
+
+# Initialize the image restoration toolkit
+tool_engine = RestorationToolkit(score_weight=[0,0,0,0,0])
+# Load the LLaVA model in half precision to reduce memory usage
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True, 
+).to(0)
+processor = AutoProcessor.from_pretrained(model_id)
+
+print("Loading tool engine...")
+
+def parse_llm_response(response):
+    """
+    Parse the LLM response to extract reason and answer sections
+    
+    Args:
+        response (str): The raw response from the LLM
+        
+    Returns:
+        tuple: (reason, answer) extracted from the response
+    """
+    reason_match = re.search(r'<reason>(.*?)</reason>', response, re.DOTALL)
+    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
+    
+    reason = reason_match.group(1).strip() if reason_match else "No reasoning provided"
+    answer = answer_match.group(1).strip() if answer_match else "No answer provided"
+    
+    return reason, answer
+
+def extract_models_from_answer(answer):
+    """
+    Extract model names from the answer string using regex
+    
+    Args:
+        answer (str): The answer string containing model recommendations
+        
+    Returns:
+        list: List of extracted model names
+    """
+    # Pattern to match [type:xxx]:(model:xxx)
+    pattern = r'\[type:[^\]]+\]:\(model:([^)]+)\)'
+    models = re.findall(pattern, answer)
+    return models
+
+def beautify_recommended_actions(answer, models):
+    """
+    Format the LLM's recommendations in a more visually appealing way
+    
+    Args:
+        answer (str): The raw answer from LLM
+        models (list): List of extracted model names
+        
+    Returns:
+        str: Beautified display of recommendations
+    """
+    
+    # Task type to emoji mapping for visual enhancement
+    task_icons = {
+        'denoise': '🧹',
+        'lighten': '💡', 
+        'derain': '🌧️',
+        'defog': '🌫️',
+        'desnow': '❄️',
+        'super_resolution': '🔍'
+    }
+    
+    # Parse the answer to extract tasks and models
+    pattern = r'\[type:([^\]]+)\]:\(model:([^)]+)\)'
+    matches = re.findall(pattern, answer)
+    
+    if not matches:
+        return f"**🎯 Recommended Actions:**\n\n{answer}\n\n**Extracted Models:** {', '.join(models) if models else 'None'}"
+    
+    # Create beautified display
+    beautified = "**🎯 Recommended Actions:**\n"
+    beautified += "> "
+    
+    # Create horizontal flow of actions
+    action_parts = []
+    for task_type, model_name in matches:
+        task_type = task_type.strip()
+        model_name = model_name.strip()
+        
+        # Get icon for task type
+        icon = task_icons.get(task_type, '🔧')
+        
+        # Format task name (capitalize and replace underscores)
+        task_display = task_type.title().replace('_', ' ')
+        
+        # Create action part: icon + task + model
+        action_part = f"{icon} {task_display}：`{model_name}`"
+        action_parts.append(action_part)
+    
+    # Join with arrows to show sequence
+    beautified += " ➡ ".join(action_parts) + "\n\n"
+    
+    # Add summary information
+    beautified += f"**📋 Processing Pipeline:** {len(matches)} steps\n"
+    beautified += f"**🛠️ Models to use:** {' → '.join(models)}"
+    
+    return beautified
+
+def resize_image_to_original(processed_image_path, original_size):
+    """
+    Resize processed image back to original dimensions
+    
+    Args:
+        processed_image_path (str): Path to the processed image
+        original_size (tuple): Original image dimensions (width, height)
+        
+    Returns:
+        str: Path to the resized image
+    """
+    if processed_image_path and os.path.exists(processed_image_path):
+        img = Image.open(processed_image_path)
+        img_resized = img.resize(original_size, Image.Resampling.LANCZOS)
+        
+        # Save resized image
+        output_path = os.path.join('temp_outputs', 'final_result.png')
+        img_resized.save(output_path)
+        return output_path
+    return processed_image_path
+
+def get_llm_response_streaming(image_path):
+    """
+    Get streaming response from LLM for image analysis
+    
+    Args:
+        image_path (str): Path to the input image
+        
+    Returns:
+        TextIteratorStreamer: A streamer object to yield tokens
+    """
+    # Select random prompt from the templates
+    instruction = prompts_query2[random.randint(0, len(prompts_query2)-1)]
+    
+    # Format the prompt with image for multimodal input
+    prompt = (f"<|start_header_id|>user<|end_header_id|>\n\n<image>\n{instruction}<|eot_id|>"
+              "<|start_header_id|>assistant<|end_header_id|>\n\n")
+    
+    # Load and process image
+    raw_image = Image.open(image_path)
+    inputs = processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
+    
+    # Setup streaming for token-by-token generation
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    # Generate response in a separate thread to avoid blocking
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=400,
+        do_sample=False
+    )
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    return streamer
+
+def process_image_with_tools(image_path, models, original_size):
+    """
+    Process image using the tool engine and restore to original size
+    
+    Args:
+        image_path (str): Path to the input image
+        models (list): List of models to apply
+        original_size (tuple): Original image dimensions
+        
+    Returns:
+        str: Path to the final processed image
+    """
+    if not models:
+        return None
+    
+    # Create output directory
+    os.makedirs('temp_outputs', exist_ok=True)
+    
+    # Process the image with selected models
+    res = tool_engine.process_image(models, image_path, 'temp_outputs')
+    
+    # Resize back to original dimensions
+    final_result = resize_image_to_original(res['output_path'], original_size)
+    
+    return final_result
+
+def process_full_pipeline(image):
+    """
+    Main processing pipeline with streaming UI updates
+    
+    Args:
+        image (str): Path to the input image
+        
+    Yields:
+        tuple: (chat_history, processed_image) for Gradio UI updates
+    """
+    if image is None:
+        return [], None
+    
+    try:
+        # Get original image size for later restoration
+        original_img = Image.open(image)
+        original_size = original_img.size
+        
+        # Initialize chat history for UI
+        chat_history = [("Image uploaded for analysis", None)]
+        
+        # Step 1: Get streaming LLM response
+        streamer = get_llm_response_streaming(image)
+        
+        # Stream the response to UI with real-time updates
+        full_response = ""
+        in_reason = False
+        in_answer = False
+        reason_displayed = False
+        answer_displayed = False
+        reasoning_added = False  # Track if reasoning entry was added
+        
+        for new_text in streamer:
+            full_response += new_text
+            
+            # Check if we're entering reason section or if we need to start showing content
+            if ('<reason>' in full_response and not in_reason and not reason_displayed) or (not reasoning_added and not in_reason and not reason_displayed):
+                in_reason = True
+                reasoning_added = True
+                
+                if '<reason>' in full_response:
+                    # Extract content after <reason>
+                    reason_start = full_response.find('<reason>') + len('<reason>')
+                    reason_content = full_response[reason_start:].strip()
+                else:
+                    # Show all content as reasoning if no tag yet
+                    reason_content = full_response.strip()
+                
+                # Add reasoning to chat history
+                chat_history.append((None, f"**🤔 Analysis & Reasoning:**\n\n{reason_content}"))
+                yield chat_history, None
+            
+            # If we're in reason section, update content
+            elif in_reason and not reason_displayed:
+                # Check if reason section is complete
+                if '</reason>' in full_response:
+                    # Extract complete reason content
+                    reason_start = full_response.find('<reason>') + len('<reason>')
+                    reason_end = full_response.find('</reason>')
+                    reason_content = full_response[reason_start:reason_end].strip()
+                    
+                    # Update chat history with complete reason
+                    chat_history[1] = (None, f"**🤔 Analysis & Reasoning:**\n\n{reason_content}")
+                    reason_displayed = True
+                    in_reason = False
+                    yield chat_history, None
+                else:
+                    # Continue streaming reason content
+                    if '<reason>' in full_response:
+                        reason_start = full_response.find('<reason>') + len('<reason>')
+                        reason_content = full_response[reason_start:].strip()
+                    else:
+                        reason_content = full_response.strip()
+                    
+                    # Update chat history with partial reason
+                    chat_history[1] = (None, f"**🤔 Analysis & Reasoning:**\n\n{reason_content}")
+                    yield chat_history, None
+            
+            # Check if we're entering answer section
+            elif '<answer>' in full_response and not in_answer and not answer_displayed and reason_displayed:
+                in_answer = True
+                # Extract content after <answer>
+                answer_start = full_response.find('<answer>') + len('<answer>')
+                answer_content = full_response[answer_start:]
+                
+                # Add partial answer to chat history
+                models = extract_models_from_answer(answer_content)
+                beautified = beautify_recommended_actions(answer_content, models)
+                chat_history.append((None, beautified))
+                yield chat_history, None
+            
+            # If we're in answer section, update content
+            elif in_answer and not answer_displayed:
+                # Check if answer section is complete
+                if '</answer>' in full_response:
+                    # Extract complete answer content
+                    answer_start = full_response.find('<answer>') + len('<answer>')
+                    answer_end = full_response.find('</answer>')
+                    answer_content = full_response[answer_start:answer_end].strip()
+                    
+                    # Parse and process final answer
+                    models = extract_models_from_answer(answer_content)
+                    beautified = beautify_recommended_actions(answer_content, models)
+                    chat_history[2] = (None, beautified)
+                    answer_displayed = True
+                    in_answer = False
+                    yield chat_history, None
+                    
+                    # Process image with tools
+                    if models:
+                        chat_history.append((None, "**🔄 Processing image...**"))
+                        yield chat_history, None
+                        
+                        processed_image = process_image_with_tools(image, models, original_size)
+                        chat_history[-1] = (None, "**✅ Processing Complete!**")
+                        yield chat_history, processed_image
+                        return
+                    else:
+                        chat_history.append((None, "**❌ No valid models found in the response**"))
+                        yield chat_history, None
+                        return
+                else:
+                    # Continue streaming answer content
+                    answer_start = full_response.find('<answer>') + len('<answer>')
+                    answer_content = full_response[answer_start:].strip()
+                    
+                    # Update chat history with partial answer
+                    models = extract_models_from_answer(answer_content)
+                    beautified = beautify_recommended_actions(answer_content, models)
+                    chat_history[2] = (None, beautified)
+                    yield chat_history, None
+        
+        # Fallback if streaming completes without proper tags
+        if not answer_displayed:
+            reason, answer = parse_llm_response(full_response)
+            models = extract_models_from_answer(answer)
+            
+            chat_history = [
+                ("Image uploaded for analysis", None),
+                (None, f"**🤔 Analysis & Reasoning:**\n\n{reason}"),
+                (None, beautify_recommended_actions(answer, models))
+            ]
+            
+            if models:
+                chat_history.append((None, "**🔄 Processing image...**"))
+                yield chat_history, None
+                
+                processed_image = process_image_with_tools(image, models, original_size)
+                chat_history[-1] = (None, "**✅ Processing Complete!**")
+                yield chat_history, processed_image
+            else:
+                chat_history.append((None, "**❌ No valid models found in the response**"))
+                yield chat_history, None
+            
+    except Exception as e:
+        error_msg = f"Error: {str(e)}"
+        chat_history = [
+            ("Image uploaded for analysis", None),
+            (None, f"**❌ Error occurred:**\n\n{error_msg}")
+        ]
+        yield chat_history, None
+
+# Create Gradio interface
+def create_interface():
+    """
+    Create and configure the Gradio web interface
+    
+    Returns:
+        gr.Blocks: Configured Gradio interface
+    """
+    with gr.Blocks(title="JarvisIR: Elevating Autonomous Driving Perception with Intelligent Image Restoration", theme=gr.themes.Soft()) as demo:
+        # Header with logo and title
+        gr.Markdown("""
+        # <img src="https://cvpr2025-jarvisir.github.io/imgs/icon.png" width="32" height="32" style="display: inline-block; vertical-align: middle; transform: translateY(-2px); margin-right: 1px;"/> JarvisIR: Elevating Autonomous Driving Perception with Intelligent Image Restoration
+        
+        Upload an image and let JarvisIR analyze its degradation and recommend the best restoration tools!
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Input image upload component
+                input_image = gr.Image(
+                    type="filepath", 
+                    label="📸 Upload Your Image",
+                    height=400
+                )
+                
+                # Process button
+                process_btn = gr.Button(
+                    "🚀 Analyze & Process", 
+                    variant="primary",
+                    size="lg"
+                )
+                
+            with gr.Column(scale=1):
+                # Chat interface to show analysis
+                chatbot = gr.Chatbot(
+                    label="💬 AI Analysis Chat",
+                    height=400,
+                    show_label=True,
+                    bubble_full_width=False
+                )
+        
+        with gr.Row():
+            # Output image display
+            output_image = gr.Image(
+                label="✨ Processed Result", 
+                height=300
+            )
+        
+        # Connect event handler for the process button
+        process_btn.click(
+            fn=process_full_pipeline,
+            inputs=[input_image],
+            outputs=[chatbot, output_image]
+        )
+        
+        # Instructions section
+        gr.Markdown("### 📝 Instructions:")
+        gr.Markdown("""
+        1. **Upload an image** that needs restoration (blurry, dark, noisy, etc.)
+        2. **Click 'Analyze & Process'** to let AI analyze the image
+        3. **View the chat** to see AI's reasoning and recommendations in real-time
+        4. **Check the result** - processed image restored to original dimensions
+        """)
+    
+    return demo
+
+if __name__ == "__main__":
+    print("Starting Image Restoration Assistant...")
+    demo = create_interface()
+    # Launch the Gradio app on specified host and port
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7866,
+        share=False
+    )

JarvisIR/docs/gradio_demo.md

@@ -0,0 +1,36 @@
+# Gradio Demo Guide
+
+This guide provides instructions on how to run the Gradio demo for JarvisIR.
+
+## Environment Setup
+
+Please follow the environment setup instructions in the [SFT Training Guide](./sft_training.md#environment-setup) to create the conda environment and install the necessary dependencies. The same environment can be used for running the Gradio demo.
+
+Make sure you have activated the conda environment:
+```bash
+conda activate sft_jarvisir
+```
+
+
+## Download Preview Weights
+
+To run the Gradio demo, you need to download the preview weights and place them in the correct location:
+
+1. Download the JarvisIR preview weights from Hugging Face
+2. Create the weights directory (if it doesn't exist):
+   ```bash
+   mkdir -p checkpoints/jarvisir-preview/
+   ```
+3. Place the downloaded weight files in the `checkpoints/jarvisir-preview/` directory
+
+
+
+## Running the Demo
+
+Once the environment is set up and activated, you can run the Gradio demo with the following command from the root directory of the project:
+
+```bash
+python gradio_demo.py
+```
+
+This will launch a web interface where you can interact with the model.

JarvisIR/docs/sft_training.md

@@ -0,0 +1,84 @@
+# SFT (Supervised Fine-Tuning) Training Guide
+
+This guide provides step-by-step instructions for setting up the environment and performing supervised fine-tuning using XTuner framework.
+
+## Environment Setup
+
+### 1. Create and Activate Virtual Environment
+
+```bash
+conda create -n sft_jarvisir python=3.10
+conda activate sft_jarvisir
+```
+
+### 2. Install Dependencies
+
+```bash
+pip install -r requirements.txt
+cd package/agent_tools/Retinexformer &&  python3 setup.py develop --no_cuda_ext && cd ..
+cd RIDCP &&  python3 setup.py develop && cd ../..
+pip install -e .
+
+cd ../dependences/BasicSR
+pip install -e .
+
+cd ../src/sft/xtuner
+
+# Install required packages
+pip install -r requirements.txt
+
+# Install XTuner with all features
+pip install -e '.[all]'
+```
+
+### 3. Verify Installation
+
+```bash
+# Check XTuner installation
+xtuner version
+
+# Verify GPU availability
+python -c "import torch; print(f'CUDA Available: {torch.cuda.is_available()}')"
+python -c "import torch; print(f'GPU Count: {torch.cuda.device_count()}')"
+```
+
+## Model Weights Setup
+
+### 1. Create Model Directory Structure
+
+```bash
+# Create directories for models and datasets
+cd ..
+mkdir -p base_models
+```
+
+### 2. Download Base Model Weights
+
+Download the following models to the base_models folder:
+<table>
+    <tr>
+        <th>Model</th><th>Model Weights</th>
+    </tr>
+    <tr style="border-top: 2px solid">
+        <td>openai/clip-vit-large-patch14-336</td><td><a href="https://huggingface.co/openai/clip-vit-large-patch14-336"> 🤗 HuggingFace</a></td>
+    </tr>
+    <tr style="border-top: 2px solid">
+        <td>xtuner/llava-llama-3-8b</td><td><a href="https://huggingface.co/xtuner/llava-llama-3-8b"> 🤗 HuggingFace</a></td>
+    </tr>
+    <tr style="border-top: 2px solid">
+        <td>xtuner/llava-llama-3-8b-pretrain</td><td><a href="https://huggingface.co/xtuner/llava-llama-3-8b-pretrain"> 🤗 HuggingFace</a></td>
+    </tr>
+</table>
+
+### 3. Training with xtuner
+- Modify the model and data paths in the **llava_8b_full.py** file. 
+- Run the following command to perform sft. 
+```bash
+NPROC_PER_NODE=${GPU_NUM} xtuner train llava_8b_full.py --deepspeed deepspeed_zero2
+```
+
+## Acknowledgements
+We would like to thank the [XTuner](https://github.com/InternLM/xtuner) team for open-sourcing such an excellent framework. For more fine-tuning methods, please refer to their official documentation and code repository.
+
+
+

JarvisIR/package/README.md

@@ -0,0 +1,155 @@
+# Image Restoration Expert Toolkit
+
+JarvisIR integrates diverse expert tools for image restoration, targeting problems like low-light conditions, rain/snow, haze, resolution enhancement, and noise. We've optimized our toolkit for both efficiency and performance, with some variations from the original paper.
+
+## Key Features
+
+- Unified framework for multiple expert models.
+- Simple API for testing each expert model.
+- Quality assessment capabilities through reward functions
+- Modular design for easy extension with new models
+
+
+## Expert Model List
+| Task | Tools | Model Description |
+|---------|---------|------|
+| **Super-resolution & Deblurring & Artifact removal** | Real-ESRGAN | Fast GAN for super-resolution, deblurring, and artifact removal. |
+| **Denoising** | SCUNet | Hybrid UNet-based model combining convolution and transformer blocks, designed for robust denoising. |
+| **Deraining** | UDR-S2Former | An uncertainty-aware transformer model for rain streak removal. |
+| | Img2img-turbo-rain | Efficient model based on SD-turbo, designed for fast and effective rain removal in real-world images. |
+| **Raindrop removal** | IDT | Transformer-based model for de-raining and raindrop removal. |
+| **Dehazing** | RIDCP | Efficient dehazing model utilizing high-quality codebook priors |
+| | KANet | Efficient dehazing network using a localization-and-removal pipeline. |
+| **Desnowing** | Img2img-turbo-snow | Efficient model for removing snow artifacts while preserving natural scene details. |
+| | Snowmaster | Real-world Image Desnowing via MLLM with Multi-Model Feedback Optimization |
+| **Low-light enhancement** | Retinexformer | One-stage Retinex-based Transformer for Low-light Image Enhancement |
+| | HVICIDNet | Lightweight transformer for low-light and exposure correction |
+| | LightenDiff | Diffusion-based framework for low-light enhancement |
+
+## Usage
+
+### Basic Usage
+
+```python
+from agent_tools import RestorationToolkit
+
+# Initialize the toolkit
+toolkit = RestorationToolkit(device='cuda')
+
+# Process an image with a sequence of models
+result = toolkit.process_image(
+    tools=['scunet', 'real_esrgan'],  # Models to apply in sequence
+    img_path='path/to/input.jpg',
+    output_dir='path/to/output'
+)
+
+# Access the result
+output_path = result['output_path']
+quality_score = result['score']
+```
+
+### Running the Test API Server
+
+```python
+from agent_tools import start_server
+
+# Start the server with default models
+start_server(host='0.0.0.0', port=5010)
+```
+
+Or use the API with curl:
+
+```bash
+curl -X POST http://localhost:5010/process_image \
+  -H "Content-Type: application/json" \
+  -d '{"img_path": "path/to/image.jpg", "models": ["scunet", "real_esrgan"]}'
+```
+
+## Cite
+
+```
+# Real-esrgan
+@inproceedings{wang2021real,
+  title={Real-esrgan: Training real-world blind super-resolution with pure synthetic data},
+  author={Wang, Xintao and Xie, Liangbin and Dong, Chao and Shan, Ying},
+  booktitle={Proceedings of the IEEE/CVF international conference on computer vision},
+  pages={1905--1914},
+  year={2021}
+}
+# img2img-turbo
+@article{parmar2024one,
+  title={One-step image translation with text-to-image models},
+  author={Parmar, Gaurav and Park, Taesung and Narasimhan, Srinivasa and Zhu, Jun-Yan},
+  journal={arXiv preprint arXiv:2403.12036},
+  year={2024}
+}
+# Lightendiffusion
+@article{jiang2024lightendiffusion,
+  title={Lightendiffusion: Unsupervised low-light image enhancement with latent-retinex diffusion models},
+  author={Jiang, Hai and Luo, Ao and Liu, Xiaohong and Han, Songchen and Liu, Shuaicheng},
+  journal={arXiv preprint arXiv:2407.08939},
+  year={2024}
+}
+# KANet
+@article{feng2024advancing,
+  title={Advancing real-world image dehazing: perspective, modules, and training},
+  author={Feng, Yuxin and Ma, Long and Meng, Xiaozhe and Zhou, Fan and Liu, Risheng and Su, Zhuo},
+  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
+  year={2024},
+  publisher={IEEE}
+}
+# IDT
+@article{xiao2022image,
+  title={Image de-raining transformer},
+  author={Xiao, Jie and Fu, Xueyang and Liu, Aiping and Wu, Feng and Zha, Zheng-Jun},
+  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
+  volume={45},
+  number={11},
+  pages={12978--12995},
+  year={2022},
+  publisher={IEEE}
+}
+# SCUNet
+@article{zhang2023practical,
+   author = {Zhang, Kai and Li, Yawei and Liang, Jingyun and Cao, Jiezhang and Zhang, Yulun and Tang, Hao and Fan, Deng-Ping and Timofte, Radu and Gool, Luc Van},
+   title = {Practical Blind Image Denoising via Swin-Conv-UNet and Data Synthesis},
+   journal = {Machine Intelligence Research},
+   DOI = {10.1007/s11633-023-1466-0},
+   url = {https://doi.org/10.1007/s11633-023-1466-0},
+   volume={20},
+   number={6},
+   pages={822--836},
+   year={2023},
+   publisher={Springer}
+}
+# S2Former
+@inproceedings{chen2023sparse,
+  title={Sparse sampling transformer with uncertainty-driven ranking for unified removal of raindrops and rain streaks},
+  author={Chen, Sixiang and Ye, Tian and Bai, Jinbin and Chen, Erkang and Shi, Jun and Zhu, Lei},
+  booktitle={Proceedings of the IEEE/CVF international conference on computer vision},
+  pages={13106--13117},
+  year={2023}
+}
+# RIDCP
+@inproceedings{wu2023ridcp,
+  title={Ridcp: Revitalizing real image dehazing via high-quality codebook priors},
+  author={Wu, Rui-Qi and Duan, Zheng-Peng and Guo, Chun-Le and Chai, Zhi and Li, Chongyi},
+  booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition},
+  pages={22282--22291},
+  year={2023}
+}
+# HVI-CIDNet 
+@article{yan2024you,
+  title={You only need one color space: An efficient network for low-light image enhancement},
+  author={Yan, Qingsen and Feng, Yixu and Zhang, Cheng and Wang, Pei and Wu, Peng and Dong, Wei and Sun, Jinqiu and Zhang, Yanning},
+  journal={arXiv preprint arXiv:2402.05809},
+  year={2024}
+}
+
+# Retinexformer
+@inproceedings{retinexformer,
+  title={Retinexformer: One-stage Retinex-based Transformer for Low-light Image Enhancement},
+  author={Yuanhao Cai and Hao Bian and Jing Lin and Haoqian Wang and Radu Timofte and Yulun Zhang},
+  booktitle={ICCV},
+  year={2023}
+}

JarvisIR/package/agent_tools.egg-info/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.4
+Name: agent_tools
+Version: 1.0.0
+Summary: Agent Tools
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Requires-Python: >=3.8

JarvisIR/package/agent_tools.egg-info/SOURCES.txt

@@ -0,0 +1,192 @@
+README.md
+pyproject.toml
+agent_tools/__init__.py
+agent_tools/iqa_reward.py
+agent_tools/restoration_toolkit.py
+agent_tools/tool_testing_api.py
+agent_tools.egg-info/PKG-INFO
+agent_tools.egg-info/SOURCES.txt
+agent_tools.egg-info/dependency_links.txt
+agent_tools.egg-info/top_level.txt
+agent_tools/ESRGAN/__init__.py
+agent_tools/ESRGAN/inference.py
+agent_tools/ESRGAN/realesrgan/__init__.py
+agent_tools/ESRGAN/realesrgan/train.py
+agent_tools/ESRGAN/realesrgan/utils.py
+agent_tools/ESRGAN/realesrgan/archs/__init__.py
+agent_tools/ESRGAN/realesrgan/archs/discriminator_arch.py
+agent_tools/ESRGAN/realesrgan/archs/srvgg_arch.py
+agent_tools/ESRGAN/realesrgan/data/__init__.py
+agent_tools/ESRGAN/realesrgan/data/realesrgan_dataset.py
+agent_tools/ESRGAN/realesrgan/data/realesrgan_paired_dataset.py
+agent_tools/ESRGAN/realesrgan/models/__init__.py
+agent_tools/ESRGAN/realesrgan/models/realesrgan_model.py
+agent_tools/ESRGAN/realesrgan/models/realesrnet_model.py
+agent_tools/HVICIDNet/inference.py
+agent_tools/HVICIDNet/mods.py
+agent_tools/HVICIDNet/wavelet.py
+agent_tools/HVICIDNet/loss/loss_utils.py
+agent_tools/HVICIDNet/loss/losses.py
+agent_tools/HVICIDNet/loss/niqe_utils.py
+agent_tools/HVICIDNet/loss/vgg_arch.py
+agent_tools/HVICIDNet/net/CIDNet.py
+agent_tools/HVICIDNet/net/HVI_transform.py
+agent_tools/HVICIDNet/net/LCA.py
+agent_tools/HVICIDNet/net/transformer_utils.py
+agent_tools/IDT/__init__.py
+agent_tools/IDT/inference.py
+agent_tools/IDT/analyse/cal_rf_bf.py
+agent_tools/IDT/models/ICRA.py
+agent_tools/IDT/models/IDT.py
+agent_tools/IDT/models/Uformer.py
+agent_tools/IDT/models/__init__.py
+agent_tools/IDT/models/atgan.py
+agent_tools/IDT/models/ddm.py
+agent_tools/IDT/models/onego_genotypes_searched.py
+agent_tools/IDT/models/onego_ops_derain.py
+agent_tools/IDT/models/onego_se_nets.py
+agent_tools/IDT/models/onego_train_model.py
+agent_tools/IDT/models/restoration.py
+agent_tools/IDT/models/restormer.py
+agent_tools/IDT/models/transformer2d.py
+agent_tools/IDT/models/unet.py
+agent_tools/IDT/utils/__init__.py
+agent_tools/IDT/utils/logging.py
+agent_tools/IDT/utils/optimize.py
+agent_tools/IDT/utils/sampling.py
+agent_tools/KANet/LD_model1.py
+agent_tools/KANet/base_networks.py
+agent_tools/KANet/deconv.py
+agent_tools/KANet/inference.py
+agent_tools/KANet/transweather_model.py
+agent_tools/LightenDiffusion/inference.py
+agent_tools/LightenDiffusion/models/__init__.py
+agent_tools/LightenDiffusion/models/ddm.py
+agent_tools/LightenDiffusion/models/decom.py
+agent_tools/LightenDiffusion/models/restoration.py
+agent_tools/LightenDiffusion/models/unet.py
+agent_tools/LightenDiffusion/utils/__init__.py
+agent_tools/LightenDiffusion/utils/logging.py
+agent_tools/LightenDiffusion/utils/optimize.py
+agent_tools/LightenDiffusion/utils/sampling.py
+agent_tools/RIDCP/__init__.py
+agent_tools/RIDCP/inference.py
+agent_tools/RIDCP/setup.py
+agent_tools/RIDCP/basicsr_ridcp/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/active_codebook.py
+agent_tools/RIDCP/basicsr_ridcp/test.py
+agent_tools/RIDCP/basicsr_ridcp/train.py
+agent_tools/RIDCP/basicsr_ridcp/archs/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/archs/arch_util.py
+agent_tools/RIDCP/basicsr_ridcp/archs/dehaze_vq_weight_arch.py
+agent_tools/RIDCP/basicsr_ridcp/archs/discriminator_arch.py
+agent_tools/RIDCP/basicsr_ridcp/archs/network_swinir.py
+agent_tools/RIDCP/basicsr_ridcp/archs/ridcp_utils.py
+agent_tools/RIDCP/basicsr_ridcp/archs/vgg_arch.py
+agent_tools/RIDCP/basicsr_ridcp/data/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/data/data_sampler.py
+agent_tools/RIDCP/basicsr_ridcp/data/data_util.py
+agent_tools/RIDCP/basicsr_ridcp/data/haze_online_dataset.py
+agent_tools/RIDCP/basicsr_ridcp/data/prefetch_dataloader.py
+agent_tools/RIDCP/basicsr_ridcp/data/transforms.py
+agent_tools/RIDCP/basicsr_ridcp/losses/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/losses/loss_util.py
+agent_tools/RIDCP/basicsr_ridcp/losses/losses.py
+agent_tools/RIDCP/basicsr_ridcp/metrics/calculate_fid_folder.py
+agent_tools/RIDCP/basicsr_ridcp/metrics/calculate_fid_stats_from_datasets.py
+agent_tools/RIDCP/basicsr_ridcp/metrics/calculate_lpips.py
+agent_tools/RIDCP/basicsr_ridcp/metrics/calculate_niqe.py
+agent_tools/RIDCP/basicsr_ridcp/metrics/calculate_psnr_ssim.py
+agent_tools/RIDCP/basicsr_ridcp/metrics/calculate_stylegan2_fid.py
+agent_tools/RIDCP/basicsr_ridcp/models/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/models/base_model.py
+agent_tools/RIDCP/basicsr_ridcp/models/dehaze_vq_model.py
+agent_tools/RIDCP/basicsr_ridcp/models/lr_scheduler.py
+agent_tools/RIDCP/basicsr_ridcp/ops/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/ops/dcn/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/ops/dcn/deform_conv.py
+agent_tools/RIDCP/basicsr_ridcp/utils/__init__.py
+agent_tools/RIDCP/basicsr_ridcp/utils/diffjpeg.py
+agent_tools/RIDCP/basicsr_ridcp/utils/dist_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/download_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/face_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/file_client.py
+agent_tools/RIDCP/basicsr_ridcp/utils/flow_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/img_process_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/img_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/lmdb_util.py
+agent_tools/RIDCP/basicsr_ridcp/utils/logger.py
+agent_tools/RIDCP/basicsr_ridcp/utils/matlab_functions.py
+agent_tools/RIDCP/basicsr_ridcp/utils/misc.py
+agent_tools/RIDCP/basicsr_ridcp/utils/options.py
+agent_tools/RIDCP/basicsr_ridcp/utils/registry.py
+agent_tools/RIDCP/utils/utils_image.py
+agent_tools/Retinexformer/__init__.py
+agent_tools/Retinexformer/inference.py
+agent_tools/Retinexformer/setup.py
+agent_tools/Retinexformer/Enhancement/utils.py
+agent_tools/Retinexformer/basicsr_retinexformer/test.py
+agent_tools/Retinexformer/basicsr_retinexformer/train.py
+agent_tools/Retinexformer/basicsr_retinexformer/version.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/SDSD_image_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/SID_image_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/SMID_image_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/__init__.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/data_sampler.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/data_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/ffhq_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/paired_image_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/prefetch_dataloader.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/reds_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/single_image_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/transforms.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/util.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/video_test_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/data/vimeo90k_dataset.py
+agent_tools/Retinexformer/basicsr_retinexformer/metrics/__init__.py
+agent_tools/Retinexformer/basicsr_retinexformer/metrics/fid.py
+agent_tools/Retinexformer/basicsr_retinexformer/metrics/metric_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/metrics/niqe.py
+agent_tools/Retinexformer/basicsr_retinexformer/metrics/psnr_ssim.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/__init__.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/base_model.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/image_restoration_model.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/lr_scheduler.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/archs/MST_Plus_Plus_arch.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/archs/RetinexFormer_arch.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/archs/__init__.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/archs/arch_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/archs/layers.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/losses/__init__.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/losses/loss_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/models/losses/losses.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/__init__.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/bundle_submissions.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/create_lmdb.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/dist_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/download_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/face_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/file_client.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/flow_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/img_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/lmdb_util.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/logger.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/matlab_functions.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/misc.py
+agent_tools/Retinexformer/basicsr_retinexformer/utils/options.py
+agent_tools/S2Former/UDR_S2Former.py
+agent_tools/S2Former/base_net_snow.py
+agent_tools/S2Former/condconv.py
+agent_tools/S2Former/inference.py
+agent_tools/SCUNet/__init__.py
+agent_tools/SCUNet/inference.py
+agent_tools/SCUNet/models/network_scunet.py
+agent_tools/SCUNet/utils/utils_image.py
+agent_tools/SnowMaster/inference.py
+agent_tools/SnowMaster/nafnet.py
+agent_tools/SnowMaster/nafnet_utils.py
+agent_tools/img2img_turbo/__init__.py
+agent_tools/img2img_turbo/inference.py
+agent_tools/img2img_turbo/src/cyclegan_turbo.py
+agent_tools/img2img_turbo/src/model.py
+agent_tools/img2img_turbo/src/my_utils/training_utils.py

JarvisIR/package/agent_tools.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

JarvisIR/package/agent_tools.egg-info/top_level.txt

@@ -0,0 +1 @@
+agent_tools

JarvisIR/package/agent_tools/ESRGAN/inference.py

@@ -0,0 +1,49 @@
+import cv2
+import os
+from basicsr.archs.rrdbnet_arch import RRDBNet
+
+from .realesrgan import RealESRGANer
+
+
+def load_esrgan_model(model_path, device):
+    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+    netscale = 4
+
+    
+    # use dni to control the denoise strength
+    dni_weight = None
+
+
+    # restorer
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=0,
+        tile_pad=10,
+        pre_pad=0,
+        half=True,
+        device=device)
+    return upsampler
+
+def esrgan_predict(upsampler, input_image, output_dir, device,):
+    
+    # determine models according to model names
+    outscale = 4 # the final upsampling scale
+    imgname, extension = os.path.splitext(os.path.basename(input_image))
+
+    img = cv2.imread(input_image, cv2.IMREAD_COLOR)
+    h, w, _ = img.shape
+    try:
+        output, _ = upsampler.enhance(img, outscale=outscale)
+    except RuntimeError as error:
+        print('Error', error)
+        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
+    # resize back to the original resolution
+    output = cv2.resize(output, (w, h), interpolation=cv2.INTER_CUBIC)
+    
+    save_path = os.path.join(output_dir, f'{imgname}.png')
+    cv2.imwrite(save_path, output)
+
+    return save_path

JarvisIR/package/agent_tools/ESRGAN/options/finetune_realesrgan_x4plus.yml

@@ -0,0 +1,188 @@
+# general settings
+name: finetune_RealESRGANx4plus_400k
+model_type: RealESRGANModel
+scale: 4
+num_gpu: auto
+manual_seed: 0
+
+# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
+# USM the ground-truth
+l1_gt_usm: True
+percep_gt_usm: True
+gan_gt_usm: False
+
+# the first degradation process
+resize_prob: [0.2, 0.7, 0.1]  # up, down, keep
+resize_range: [0.15, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 0.4
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 0.8
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 0.4
+jpeg_range2: [30, 95]
+
+gt_size: 256
+queue_size: 180
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DF2K+OST
+    type: RealESRGANDataset
+    dataroot_gt: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
+    io_backend:
+      type: disk
+
+    blur_kernel_size: 21
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob: 0.1
+    blur_sigma: [0.2, 3]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 21
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob2: 0.1
+    blur_sigma2: [0.2, 1.5]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_sinc_prob: 0.8
+
+    gt_size: 256
+    use_hflip: True
+    use_rot: False
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 5
+    batch_size_per_gpu: 12
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  # Uncomment these for validation
+  # val:
+  #   name: validation
+  #   type: PairedImageDataset
+  #   dataroot_gt: path_to_gt
+  #   dataroot_lq: path_to_lq
+  #   io_backend:
+  #     type: disk
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 64
+  num_block: 23
+  num_grow_ch: 32
+
+network_d:
+  type: UNetDiscriminatorSN
+  num_in_ch: 3
+  num_feat: 64
+  skip_connection: True
+
+# path
+path:
+  # use the pre-trained Real-ESRNet model
+  pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth
+  param_key_g: params_ema
+  strict_load_g: true
+  pretrain_network_d: experiments/pretrained_models/RealESRGAN_x4plus_netD.pth
+  param_key_d: params
+  strict_load_d: true
+  resume_state: ~
+
+# training settings
+train:
+  ema_decay: 0.999
+  optim_g:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  optim_d:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+
+  scheduler:
+    type: MultiStepLR
+    milestones: [400000]
+    gamma: 0.5
+
+  total_iter: 400000
+  warmup_iter: -1  # no warm up
+
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+  # perceptual loss (content and style losses)
+  perceptual_opt:
+    type: PerceptualLoss
+    layer_weights:
+      # before relu
+      'conv1_2': 0.1
+      'conv2_2': 0.1
+      'conv3_4': 1
+      'conv4_4': 1
+      'conv5_4': 1
+    vgg_type: vgg19
+    use_input_norm: true
+    perceptual_weight: !!float 1.0
+    style_weight: 0
+    range_norm: false
+    criterion: l1
+  # gan loss
+  gan_opt:
+    type: GANLoss
+    gan_type: vanilla
+    real_label_val: 1.0
+    fake_label_val: 0.0
+    loss_weight: !!float 1e-1
+
+  net_d_iters: 1
+  net_d_init_iters: 0
+
+# Uncomment these for validation
+# validation settings
+# val:
+#   val_freq: !!float 5e3
+#   save_img: True
+
+#   metrics:
+#     psnr: # metric name
+#       type: calculate_psnr
+#       crop_border: 4
+#       test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3
+  use_tb_logger: true
+  wandb:
+    project: ~
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

JarvisIR/package/agent_tools/ESRGAN/options/finetune_realesrgan_x4plus_pairdata.yml

@@ -0,0 +1,150 @@
+# general settings
+name: finetune_RealESRGANx4plus_400k_pairdata
+model_type: RealESRGANModel
+scale: 4
+num_gpu: auto
+manual_seed: 0
+
+# USM the ground-truth
+l1_gt_usm: True
+percep_gt_usm: True
+gan_gt_usm: False
+
+high_order_degradation: False # do not use the high-order degradation generation process
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DIV2K
+    type: RealESRGANPairedDataset
+    dataroot_gt: datasets/DF2K
+    dataroot_lq: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt
+    io_backend:
+      type: disk
+
+    gt_size: 256
+    use_hflip: True
+    use_rot: False
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 5
+    batch_size_per_gpu: 12
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  # Uncomment these for validation
+  # val:
+  #   name: validation
+  #   type: PairedImageDataset
+  #   dataroot_gt: path_to_gt
+  #   dataroot_lq: path_to_lq
+  #   io_backend:
+  #     type: disk
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 64
+  num_block: 23
+  num_grow_ch: 32
+
+network_d:
+  type: UNetDiscriminatorSN
+  num_in_ch: 3
+  num_feat: 64
+  skip_connection: True
+
+# path
+path:
+  # use the pre-trained Real-ESRNet model
+  pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth
+  param_key_g: params_ema
+  strict_load_g: true
+  pretrain_network_d: experiments/pretrained_models/RealESRGAN_x4plus_netD.pth
+  param_key_d: params
+  strict_load_d: true
+  resume_state: ~
+
+# training settings
+train:
+  ema_decay: 0.999
+  optim_g:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  optim_d:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+
+  scheduler:
+    type: MultiStepLR
+    milestones: [400000]
+    gamma: 0.5
+
+  total_iter: 400000
+  warmup_iter: -1  # no warm up
+
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+  # perceptual loss (content and style losses)
+  perceptual_opt:
+    type: PerceptualLoss
+    layer_weights:
+      # before relu
+      'conv1_2': 0.1
+      'conv2_2': 0.1
+      'conv3_4': 1
+      'conv4_4': 1
+      'conv5_4': 1
+    vgg_type: vgg19
+    use_input_norm: true
+    perceptual_weight: !!float 1.0
+    style_weight: 0
+    range_norm: false
+    criterion: l1
+  # gan loss
+  gan_opt:
+    type: GANLoss
+    gan_type: vanilla
+    real_label_val: 1.0
+    fake_label_val: 0.0
+    loss_weight: !!float 1e-1
+
+  net_d_iters: 1
+  net_d_init_iters: 0
+
+# Uncomment these for validation
+# validation settings
+# val:
+#   val_freq: !!float 5e3
+#   save_img: True
+
+#   metrics:
+#     psnr: # metric name
+#       type: calculate_psnr
+#       crop_border: 4
+#       test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3
+  use_tb_logger: true
+  wandb:
+    project: ~
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

JarvisIR/package/agent_tools/ESRGAN/options/train_realesrgan_x2plus.yml

@@ -0,0 +1,186 @@
+# general settings
+name: train_RealESRGANx2plus_400k_B12G4
+model_type: RealESRGANModel
+scale: 2
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
+manual_seed: 0
+
+# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
+# USM the ground-truth
+l1_gt_usm: True
+percep_gt_usm: True
+gan_gt_usm: False
+
+# the first degradation process
+resize_prob: [0.2, 0.7, 0.1]  # up, down, keep
+resize_range: [0.15, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 0.4
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 0.8
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 0.4
+jpeg_range2: [30, 95]
+
+gt_size: 256
+queue_size: 180
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DF2K+OST
+    type: RealESRGANDataset
+    dataroot_gt: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
+    io_backend:
+      type: disk
+
+    blur_kernel_size: 21
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob: 0.1
+    blur_sigma: [0.2, 3]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 21
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob2: 0.1
+    blur_sigma2: [0.2, 1.5]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_sinc_prob: 0.8
+
+    gt_size: 256
+    use_hflip: True
+    use_rot: False
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 5
+    batch_size_per_gpu: 12
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  # Uncomment these for validation
+  # val:
+  #   name: validation
+  #   type: PairedImageDataset
+  #   dataroot_gt: path_to_gt
+  #   dataroot_lq: path_to_lq
+  #   io_backend:
+  #     type: disk
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 64
+  num_block: 23
+  num_grow_ch: 32
+  scale: 2
+
+network_d:
+  type: UNetDiscriminatorSN
+  num_in_ch: 3
+  num_feat: 64
+  skip_connection: True
+
+# path
+path:
+  # use the pre-trained Real-ESRNet model
+  pretrain_network_g: experiments/pretrained_models/RealESRNet_x2plus.pth
+  param_key_g: params_ema
+  strict_load_g: true
+  resume_state: ~
+
+# training settings
+train:
+  ema_decay: 0.999
+  optim_g:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  optim_d:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+
+  scheduler:
+    type: MultiStepLR
+    milestones: [400000]
+    gamma: 0.5
+
+  total_iter: 400000
+  warmup_iter: -1  # no warm up
+
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+  # perceptual loss (content and style losses)
+  perceptual_opt:
+    type: PerceptualLoss
+    layer_weights:
+      # before relu
+      'conv1_2': 0.1
+      'conv2_2': 0.1
+      'conv3_4': 1
+      'conv4_4': 1
+      'conv5_4': 1
+    vgg_type: vgg19
+    use_input_norm: true
+    perceptual_weight: !!float 1.0
+    style_weight: 0
+    range_norm: false
+    criterion: l1
+  # gan loss
+  gan_opt:
+    type: GANLoss
+    gan_type: vanilla
+    real_label_val: 1.0
+    fake_label_val: 0.0
+    loss_weight: !!float 1e-1
+
+  net_d_iters: 1
+  net_d_init_iters: 0
+
+# Uncomment these for validation
+# validation settings
+# val:
+#   val_freq: !!float 5e3
+#   save_img: True
+
+#   metrics:
+#     psnr: # metric name
+#       type: calculate_psnr
+#       crop_border: 4
+#       test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3
+  use_tb_logger: true
+  wandb:
+    project: ~
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

JarvisIR/package/agent_tools/ESRGAN/options/train_realesrgan_x4plus.yml

@@ -0,0 +1,185 @@
+# general settings
+name: train_RealESRGANx4plus_400k_B12G4
+model_type: RealESRGANModel
+scale: 4
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
+manual_seed: 0
+
+# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
+# USM the ground-truth
+l1_gt_usm: True
+percep_gt_usm: True
+gan_gt_usm: False
+
+# the first degradation process
+resize_prob: [0.2, 0.7, 0.1]  # up, down, keep
+resize_range: [0.15, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 0.4
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 0.8
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 0.4
+jpeg_range2: [30, 95]
+
+gt_size: 256
+queue_size: 180
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DF2K+OST
+    type: RealESRGANDataset
+    dataroot_gt: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
+    io_backend:
+      type: disk
+
+    blur_kernel_size: 21
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob: 0.1
+    blur_sigma: [0.2, 3]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 21
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob2: 0.1
+    blur_sigma2: [0.2, 1.5]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_sinc_prob: 0.8
+
+    gt_size: 256
+    use_hflip: True
+    use_rot: False
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 5
+    batch_size_per_gpu: 12
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  # Uncomment these for validation
+  # val:
+  #   name: validation
+  #   type: PairedImageDataset
+  #   dataroot_gt: path_to_gt
+  #   dataroot_lq: path_to_lq
+  #   io_backend:
+  #     type: disk
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 64
+  num_block: 23
+  num_grow_ch: 32
+
+network_d:
+  type: UNetDiscriminatorSN
+  num_in_ch: 3
+  num_feat: 64
+  skip_connection: True
+
+# path
+path:
+  # use the pre-trained Real-ESRNet model
+  pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth
+  param_key_g: params_ema
+  strict_load_g: true
+  resume_state: ~
+
+# training settings
+train:
+  ema_decay: 0.999
+  optim_g:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  optim_d:
+    type: Adam
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+
+  scheduler:
+    type: MultiStepLR
+    milestones: [400000]
+    gamma: 0.5
+
+  total_iter: 400000
+  warmup_iter: -1  # no warm up
+
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+  # perceptual loss (content and style losses)
+  perceptual_opt:
+    type: PerceptualLoss
+    layer_weights:
+      # before relu
+      'conv1_2': 0.1
+      'conv2_2': 0.1
+      'conv3_4': 1
+      'conv4_4': 1
+      'conv5_4': 1
+    vgg_type: vgg19
+    use_input_norm: true
+    perceptual_weight: !!float 1.0
+    style_weight: 0
+    range_norm: false
+    criterion: l1
+  # gan loss
+  gan_opt:
+    type: GANLoss
+    gan_type: vanilla
+    real_label_val: 1.0
+    fake_label_val: 0.0
+    loss_weight: !!float 1e-1
+
+  net_d_iters: 1
+  net_d_init_iters: 0
+
+# Uncomment these for validation
+# validation settings
+# val:
+#   val_freq: !!float 5e3
+#   save_img: True
+
+#   metrics:
+#     psnr: # metric name
+#       type: calculate_psnr
+#       crop_border: 4
+#       test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3
+  use_tb_logger: true
+  wandb:
+    project: ~
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

JarvisIR/package/agent_tools/ESRGAN/options/train_realesrnet_x2plus.yml

@@ -0,0 +1,145 @@
+# general settings
+name: train_RealESRNetx2plus_1000k_B12G4
+model_type: RealESRNetModel
+scale: 2
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
+manual_seed: 0
+
+# ----------------- options for synthesizing training data in RealESRNetModel ----------------- #
+gt_usm: True  # USM the ground-truth
+
+# the first degradation process
+resize_prob: [0.2, 0.7, 0.1]  # up, down, keep
+resize_range: [0.15, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 0.4
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 0.8
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 0.4
+jpeg_range2: [30, 95]
+
+gt_size: 256
+queue_size: 180
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DF2K+OST
+    type: RealESRGANDataset
+    dataroot_gt: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
+    io_backend:
+      type: disk
+
+    blur_kernel_size: 21
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob: 0.1
+    blur_sigma: [0.2, 3]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 21
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob2: 0.1
+    blur_sigma2: [0.2, 1.5]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_sinc_prob: 0.8
+
+    gt_size: 256
+    use_hflip: True
+    use_rot: False
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 5
+    batch_size_per_gpu: 12
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  # Uncomment these for validation
+  # val:
+  #   name: validation
+  #   type: PairedImageDataset
+  #   dataroot_gt: path_to_gt
+  #   dataroot_lq: path_to_lq
+  #   io_backend:
+  #     type: disk
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 64
+  num_block: 23
+  num_grow_ch: 32
+  scale: 2
+
+# path
+path:
+  pretrain_network_g: experiments/pretrained_models/RealESRGAN_x4plus.pth
+  param_key_g: params_ema
+  strict_load_g: False
+  resume_state: ~
+
+# training settings
+train:
+  ema_decay: 0.999
+  optim_g:
+    type: Adam
+    lr: !!float 2e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+
+  scheduler:
+    type: MultiStepLR
+    milestones: [1000000]
+    gamma: 0.5
+
+  total_iter: 1000000
+  warmup_iter: -1  # no warm up
+
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+
+# Uncomment these for validation
+# validation settings
+# val:
+#   val_freq: !!float 5e3
+#   save_img: True
+
+#   metrics:
+#     psnr: # metric name
+#       type: calculate_psnr
+#       crop_border: 4
+#       test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3
+  use_tb_logger: true
+  wandb:
+    project: ~
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

JarvisIR/package/agent_tools/ESRGAN/options/train_realesrnet_x4plus.yml

@@ -0,0 +1,144 @@
+# general settings
+name: train_RealESRNetx4plus_1000k_B12G4
+model_type: RealESRNetModel
+scale: 4
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
+manual_seed: 0
+
+# ----------------- options for synthesizing training data in RealESRNetModel ----------------- #
+gt_usm: True  # USM the ground-truth
+
+# the first degradation process
+resize_prob: [0.2, 0.7, 0.1]  # up, down, keep
+resize_range: [0.15, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 0.4
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 0.8
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 0.4
+jpeg_range2: [30, 95]
+
+gt_size: 256
+queue_size: 180
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DF2K+OST
+    type: RealESRGANDataset
+    dataroot_gt: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
+    io_backend:
+      type: disk
+
+    blur_kernel_size: 21
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob: 0.1
+    blur_sigma: [0.2, 3]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 21
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    sinc_prob2: 0.1
+    blur_sigma2: [0.2, 1.5]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_sinc_prob: 0.8
+
+    gt_size: 256
+    use_hflip: True
+    use_rot: False
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 5
+    batch_size_per_gpu: 12
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  # Uncomment these for validation
+  # val:
+  #   name: validation
+  #   type: PairedImageDataset
+  #   dataroot_gt: path_to_gt
+  #   dataroot_lq: path_to_lq
+  #   io_backend:
+  #     type: disk
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 64
+  num_block: 23
+  num_grow_ch: 32
+
+# path
+path:
+  pretrain_network_g: experiments/pretrained_models/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth
+  param_key_g: params_ema
+  strict_load_g: true
+  resume_state: ~
+
+# training settings
+train:
+  ema_decay: 0.999
+  optim_g:
+    type: Adam
+    lr: !!float 2e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+
+  scheduler:
+    type: MultiStepLR
+    milestones: [1000000]
+    gamma: 0.5
+
+  total_iter: 1000000
+  warmup_iter: -1  # no warm up
+
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+
+# Uncomment these for validation
+# validation settings
+# val:
+#   val_freq: !!float 5e3
+#   save_img: True
+
+#   metrics:
+#     psnr: # metric name
+#       type: calculate_psnr
+#       crop_border: 4
+#       test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3
+  use_tb_logger: true
+  wandb:
+    project: ~
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

JarvisIR/package/agent_tools/ESRGAN/realesrgan/__init__.py

@@ -0,0 +1,5 @@
+# flake8: noqa
+from .archs import *
+from .data import *
+from .models import *
+from .utils import *

JarvisIR/package/agent_tools/ESRGAN/realesrgan/archs/__init__.py

@@ -0,0 +1,11 @@
+import importlib
+from basicsr.utils import scandir
+from os import path as osp
+
+# automatically scan and import arch modules for registry
+# scan all the files that end with '_arch.py' under the archs folder
+arch_folder = osp.dirname(osp.abspath(__file__))
+arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_arch.py')]
+# import all the arch modules
+from . import srvgg_arch
+from . import discriminator_arch

JarvisIR/package/agent_tools/ESRGAN/realesrgan/archs/discriminator_arch.py

@@ -0,0 +1,67 @@
+from basicsr.utils.registry import ARCH_REGISTRY
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm
+
+
+@ARCH_REGISTRY.register()
+class UNetDiscriminatorSN(nn.Module):
+    """Defines a U-Net discriminator with spectral normalization (SN)
+
+    It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    Arg:
+        num_in_ch (int): Channel number of inputs. Default: 3.
+        num_feat (int): Channel number of base intermediate features. Default: 64.
+        skip_connection (bool): Whether to use skip connections between U-Net. Default: True.
+    """
+
+    def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
+        super(UNetDiscriminatorSN, self).__init__()
+        self.skip_connection = skip_connection
+        norm = spectral_norm
+        # the first convolution
+        self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
+        # downsample
+        self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))
+        self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))
+        self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))
+        # upsample
+        self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
+        self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
+        self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))
+        # extra convolutions
+        self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
+        self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
+        self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)
+
+    def forward(self, x):
+        # downsample
+        x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
+        x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
+        x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
+        x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)
+
+        # upsample
+        x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)
+        x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)
+
+        if self.skip_connection:
+            x4 = x4 + x2
+        x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)
+        x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)
+
+        if self.skip_connection:
+            x5 = x5 + x1
+        x5 = F.interpolate(x5, scale_factor=2, mode='bilinear', align_corners=False)
+        x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)
+
+        if self.skip_connection:
+            x6 = x6 + x0
+
+        # extra convolutions
+        out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
+        out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
+        out = self.conv9(out)
+
+        return out

JarvisIR/package/agent_tools/ESRGAN/realesrgan/archs/srvgg_arch.py

@@ -0,0 +1,69 @@
+from basicsr.utils.registry import ARCH_REGISTRY
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+@ARCH_REGISTRY.register()
+class SRVGGNetCompact(nn.Module):
+    """A compact VGG-style network structure for super-resolution.
+
+    It is a compact network structure, which performs upsampling in the last layer and no convolution is
+    conducted on the HR feature space.
+
+    Args:
+        num_in_ch (int): Channel number of inputs. Default: 3.
+        num_out_ch (int): Channel number of outputs. Default: 3.
+        num_feat (int): Channel number of intermediate features. Default: 64.
+        num_conv (int): Number of convolution layers in the body network. Default: 16.
+        upscale (int): Upsampling factor. Default: 4.
+        act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu.
+    """
+
+    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'):
+        super(SRVGGNetCompact, self).__init__()
+        self.num_in_ch = num_in_ch
+        self.num_out_ch = num_out_ch
+        self.num_feat = num_feat
+        self.num_conv = num_conv
+        self.upscale = upscale
+        self.act_type = act_type
+
+        self.body = nn.ModuleList()
+        # the first conv
+        self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))
+        # the first activation
+        if act_type == 'relu':
+            activation = nn.ReLU(inplace=True)
+        elif act_type == 'prelu':
+            activation = nn.PReLU(num_parameters=num_feat)
+        elif act_type == 'leakyrelu':
+            activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+        self.body.append(activation)
+
+        # the body structure
+        for _ in range(num_conv):
+            self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))
+            # activation
+            if act_type == 'relu':
+                activation = nn.ReLU(inplace=True)
+            elif act_type == 'prelu':
+                activation = nn.PReLU(num_parameters=num_feat)
+            elif act_type == 'leakyrelu':
+                activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+            self.body.append(activation)
+
+        # the last conv
+        self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))
+        # upsample
+        self.upsampler = nn.PixelShuffle(upscale)
+
+    def forward(self, x):
+        out = x
+        for i in range(0, len(self.body)):
+            out = self.body[i](out)
+
+        out = self.upsampler(out)
+        # add the nearest upsampled image, so that the network learns the residual
+        base = F.interpolate(x, scale_factor=self.upscale, mode='nearest')
+        out += base
+        return out

JarvisIR/package/agent_tools/ESRGAN/realesrgan/data/__init__.py

@@ -0,0 +1,11 @@
+import importlib
+from basicsr.utils import scandir
+from os import path as osp
+
+# automatically scan and import dataset modules for registry
+# scan all the files that end with '_dataset.py' under the data folder
+data_folder = osp.dirname(osp.abspath(__file__))
+dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
+# import all the dataset modules
+from . import realesrgan_paired_dataset
+from . import realesrgan_dataset

JarvisIR/package/agent_tools/ESRGAN/realesrgan/data/realesrgan_dataset.py

@@ -0,0 +1,192 @@
+import cv2
+import math
+import numpy as np
+import os
+import os.path as osp
+import random
+import time
+import torch
+from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels
+from basicsr.data.transforms import augment
+from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
+from basicsr.utils.registry import DATASET_REGISTRY
+from torch.utils import data as data
+
+
+@DATASET_REGISTRY.register()
+class RealESRGANDataset(data.Dataset):
+    """Dataset used for Real-ESRGAN model:
+    Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    It loads gt (Ground-Truth) images, and augments them.
+    It also generates blur kernels and sinc kernels for generating low-quality images.
+    Note that the low-quality images are processed in tensors on GPUS for faster processing.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_gt (str): Data root path for gt.
+            meta_info (str): Path for meta information file.
+            io_backend (dict): IO backend type and other kwarg.
+            use_hflip (bool): Use horizontal flips.
+            use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
+            Please see more options in the codes.
+    """
+
+    def __init__(self, opt):
+        super(RealESRGANDataset, self).__init__()
+        self.opt = opt
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+        self.gt_folder = opt['dataroot_gt']
+
+        # file client (lmdb io backend)
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = [self.gt_folder]
+            self.io_backend_opt['client_keys'] = ['gt']
+            if not self.gt_folder.endswith('.lmdb'):
+                raise ValueError(f"'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}")
+            with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
+                self.paths = [line.split('.')[0] for line in fin]
+        else:
+            # disk backend with meta_info
+            # Each line in the meta_info describes the relative path to an image
+            with open(self.opt['meta_info']) as fin:
+                paths = [line.strip().split(' ')[0] for line in fin]
+                self.paths = [os.path.join(self.gt_folder, v) for v in paths]
+
+        # blur settings for the first degradation
+        self.blur_kernel_size = opt['blur_kernel_size']
+        self.kernel_list = opt['kernel_list']
+        self.kernel_prob = opt['kernel_prob']  # a list for each kernel probability
+        self.blur_sigma = opt['blur_sigma']
+        self.betag_range = opt['betag_range']  # betag used in generalized Gaussian blur kernels
+        self.betap_range = opt['betap_range']  # betap used in plateau blur kernels
+        self.sinc_prob = opt['sinc_prob']  # the probability for sinc filters
+
+        # blur settings for the second degradation
+        self.blur_kernel_size2 = opt['blur_kernel_size2']
+        self.kernel_list2 = opt['kernel_list2']
+        self.kernel_prob2 = opt['kernel_prob2']
+        self.blur_sigma2 = opt['blur_sigma2']
+        self.betag_range2 = opt['betag_range2']
+        self.betap_range2 = opt['betap_range2']
+        self.sinc_prob2 = opt['sinc_prob2']
+
+        # a final sinc filter
+        self.final_sinc_prob = opt['final_sinc_prob']
+
+        self.kernel_range = [2 * v + 1 for v in range(3, 11)]  # kernel size ranges from 7 to 21
+        # TODO: kernel range is now hard-coded, should be in the configure file
+        self.pulse_tensor = torch.zeros(21, 21).float()  # convolving with pulse tensor brings no blurry effect
+        self.pulse_tensor[10, 10] = 1
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        # -------------------------------- Load gt images -------------------------------- #
+        # Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32.
+        gt_path = self.paths[index]
+        # avoid errors caused by high latency in reading files
+        retry = 3
+        while retry > 0:
+            try:
+                img_bytes = self.file_client.get(gt_path, 'gt')
+            except (IOError, OSError) as e:
+                logger = get_root_logger()
+                logger.warn(f'File client error: {e}, remaining retry times: {retry - 1}')
+                # change another file to read
+                index = random.randint(0, self.__len__())
+                gt_path = self.paths[index]
+                time.sleep(1)  # sleep 1s for occasional server congestion
+            else:
+                break
+            finally:
+                retry -= 1
+        img_gt = imfrombytes(img_bytes, float32=True)
+
+        # -------------------- Do augmentation for training: flip, rotation -------------------- #
+        img_gt = augment(img_gt, self.opt['use_hflip'], self.opt['use_rot'])
+
+        # crop or pad to 400
+        # TODO: 400 is hard-coded. You may change it accordingly
+        h, w = img_gt.shape[0:2]
+        crop_pad_size = 400
+        # pad
+        if h < crop_pad_size or w < crop_pad_size:
+            pad_h = max(0, crop_pad_size - h)
+            pad_w = max(0, crop_pad_size - w)
+            img_gt = cv2.copyMakeBorder(img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
+        # crop
+        if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size:
+            h, w = img_gt.shape[0:2]
+            # randomly choose top and left coordinates
+            top = random.randint(0, h - crop_pad_size)
+            left = random.randint(0, w - crop_pad_size)
+            img_gt = img_gt[top:top + crop_pad_size, left:left + crop_pad_size, ...]
+
+        # ------------------------ Generate kernels (used in the first degradation) ------------------------ #
+        kernel_size = random.choice(self.kernel_range)
+        if np.random.uniform() < self.opt['sinc_prob']:
+            # this sinc filter setting is for kernels ranging from [7, 21]
+            if kernel_size < 13:
+                omega_c = np.random.uniform(np.pi / 3, np.pi)
+            else:
+                omega_c = np.random.uniform(np.pi / 5, np.pi)
+            kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
+        else:
+            kernel = random_mixed_kernels(
+                self.kernel_list,
+                self.kernel_prob,
+                kernel_size,
+                self.blur_sigma,
+                self.blur_sigma, [-math.pi, math.pi],
+                self.betag_range,
+                self.betap_range,
+                noise_range=None)
+        # pad kernel
+        pad_size = (21 - kernel_size) // 2
+        kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
+
+        # ------------------------ Generate kernels (used in the second degradation) ------------------------ #
+        kernel_size = random.choice(self.kernel_range)
+        if np.random.uniform() < self.opt['sinc_prob2']:
+            if kernel_size < 13:
+                omega_c = np.random.uniform(np.pi / 3, np.pi)
+            else:
+                omega_c = np.random.uniform(np.pi / 5, np.pi)
+            kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
+        else:
+            kernel2 = random_mixed_kernels(
+                self.kernel_list2,
+                self.kernel_prob2,
+                kernel_size,
+                self.blur_sigma2,
+                self.blur_sigma2, [-math.pi, math.pi],
+                self.betag_range2,
+                self.betap_range2,
+                noise_range=None)
+
+        # pad kernel
+        pad_size = (21 - kernel_size) // 2
+        kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
+
+        # ------------------------------------- the final sinc kernel ------------------------------------- #
+        if np.random.uniform() < self.opt['final_sinc_prob']:
+            kernel_size = random.choice(self.kernel_range)
+            omega_c = np.random.uniform(np.pi / 3, np.pi)
+            sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
+            sinc_kernel = torch.FloatTensor(sinc_kernel)
+        else:
+            sinc_kernel = self.pulse_tensor
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
+        kernel = torch.FloatTensor(kernel)
+        kernel2 = torch.FloatTensor(kernel2)
+
+        return_d = {'gt': img_gt, 'kernel1': kernel, 'kernel2': kernel2, 'sinc_kernel': sinc_kernel, 'gt_path': gt_path}
+        return return_d
+
+    def __len__(self):
+        return len(self.paths)

JarvisIR/package/agent_tools/ESRGAN/realesrgan/data/realesrgan_paired_dataset.py

@@ -0,0 +1,108 @@
+import os
+from basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb
+from basicsr.data.transforms import augment, paired_random_crop
+from basicsr.utils import FileClient, imfrombytes, img2tensor
+from basicsr.utils.registry import DATASET_REGISTRY
+from torch.utils import data as data
+from torchvision.transforms.functional import normalize
+
+
+@DATASET_REGISTRY.register()
+class RealESRGANPairedDataset(data.Dataset):
+    """Paired image dataset for image restoration.
+
+    Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
+
+    There are three modes:
+    1. 'lmdb': Use lmdb files.
+        If opt['io_backend'] == lmdb.
+    2. 'meta_info': Use meta information file to generate paths.
+        If opt['io_backend'] != lmdb and opt['meta_info'] is not None.
+    3. 'folder': Scan folders to generate paths.
+        The rest.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_gt (str): Data root path for gt.
+            dataroot_lq (str): Data root path for lq.
+            meta_info (str): Path for meta information file.
+            io_backend (dict): IO backend type and other kwarg.
+            filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
+                Default: '{}'.
+            gt_size (int): Cropped patched size for gt patches.
+            use_hflip (bool): Use horizontal flips.
+            use_rot (bool): Use rotation (use vertical flip and transposing h
+                and w for implementation).
+
+            scale (bool): Scale, which will be added automatically.
+            phase (str): 'train' or 'val'.
+    """
+
+    def __init__(self, opt):
+        super(RealESRGANPairedDataset, self).__init__()
+        self.opt = opt
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+        # mean and std for normalizing the input images
+        self.mean = opt['mean'] if 'mean' in opt else None
+        self.std = opt['std'] if 'std' in opt else None
+
+        self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
+        self.filename_tmpl = opt['filename_tmpl'] if 'filename_tmpl' in opt else '{}'
+
+        # file client (lmdb io backend)
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
+            self.io_backend_opt['client_keys'] = ['lq', 'gt']
+            self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
+        elif 'meta_info' in self.opt and self.opt['meta_info'] is not None:
+            # disk backend with meta_info
+            # Each line in the meta_info describes the relative path to an image
+            with open(self.opt['meta_info']) as fin:
+                paths = [line.strip() for line in fin]
+            self.paths = []
+            for path in paths:
+                gt_path, lq_path = path.split(', ')
+                gt_path = os.path.join(self.gt_folder, gt_path)
+                lq_path = os.path.join(self.lq_folder, lq_path)
+                self.paths.append(dict([('gt_path', gt_path), ('lq_path', lq_path)]))
+        else:
+            # disk backend
+            # it will scan the whole folder to get meta info
+            # it will be time-consuming for folders with too many files. It is recommended using an extra meta txt file
+            self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        scale = self.opt['scale']
+
+        # Load gt and lq images. Dimension order: HWC; channel order: BGR;
+        # image range: [0, 1], float32.
+        gt_path = self.paths[index]['gt_path']
+        img_bytes = self.file_client.get(gt_path, 'gt')
+        img_gt = imfrombytes(img_bytes, float32=True)
+        lq_path = self.paths[index]['lq_path']
+        img_bytes = self.file_client.get(lq_path, 'lq')
+        img_lq = imfrombytes(img_bytes, float32=True)
+
+        # augmentation for training
+        if self.opt['phase'] == 'train':
+            gt_size = self.opt['gt_size']
+            # random crop
+            img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
+            # flip, rotation
+            img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
+        # normalize
+        if self.mean is not None or self.std is not None:
+            normalize(img_lq, self.mean, self.std, inplace=True)
+            normalize(img_gt, self.mean, self.std, inplace=True)
+
+        return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
+
+    def __len__(self):
+        return len(self.paths)

JarvisIR/package/agent_tools/ESRGAN/realesrgan/models/__init__.py

@@ -0,0 +1,11 @@
+import importlib
+from basicsr.utils import scandir
+from os import path as osp
+
+# automatically scan and import model modules for registry
+# scan all the files that end with '_model.py' under the model folder
+model_folder = osp.dirname(osp.abspath(__file__))
+model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
+# import all the model modules
+from . import realesrgan_model
+from . import realesrnet_model

JarvisIR/package/agent_tools/ESRGAN/realesrgan/models/realesrgan_model.py

@@ -0,0 +1,258 @@
+import numpy as np
+import random
+import torch
+from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
+from basicsr.data.transforms import paired_random_crop
+from basicsr.models.srgan_model import SRGANModel
+from basicsr.utils import DiffJPEG, USMSharp
+from basicsr.utils.img_process_util import filter2D
+from basicsr.utils.registry import MODEL_REGISTRY
+from collections import OrderedDict
+from torch.nn import functional as F
+
+
+@MODEL_REGISTRY.register()
+class RealESRGANModel(SRGANModel):
+    """RealESRGAN Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    It mainly performs:
+    1. randomly synthesize LQ images in GPU tensors
+    2. optimize the networks with GAN training.
+    """
+
+    def __init__(self, opt):
+        super(RealESRGANModel, self).__init__(opt)
+        self.jpeger = DiffJPEG(differentiable=False).cuda()  # simulate JPEG compression artifacts
+        self.usm_sharpener = USMSharp().cuda()  # do usm sharpening
+        self.queue_size = opt.get('queue_size', 180)
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self):
+        """It is the training pair pool for increasing the diversity in a batch.
+
+        Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
+        batch could not have different resize scaling factors. Therefore, we employ this training pair pool
+        to increase the degradation diversity in a batch.
+        """
+        # initialize
+        b, c, h, w = self.lq.size()
+        if not hasattr(self, 'queue_lr'):
+            assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
+            self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
+            _, c, h, w = self.gt.size()
+            self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
+            self.queue_ptr = 0
+        if self.queue_ptr == self.queue_size:  # the pool is full
+            # do dequeue and enqueue
+            # shuffle
+            idx = torch.randperm(self.queue_size)
+            self.queue_lr = self.queue_lr[idx]
+            self.queue_gt = self.queue_gt[idx]
+            # get first b samples
+            lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
+            gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
+            # update the queue
+            self.queue_lr[0:b, :, :, :] = self.lq.clone()
+            self.queue_gt[0:b, :, :, :] = self.gt.clone()
+
+            self.lq = lq_dequeue
+            self.gt = gt_dequeue
+        else:
+            # only do enqueue
+            self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
+            self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
+            self.queue_ptr = self.queue_ptr + b
+
+    @torch.no_grad()
+    def feed_data(self, data):
+        """Accept data from dataloader, and then add two-order degradations to obtain LQ images.
+        """
+        if self.is_train and self.opt.get('high_order_degradation', True):
+            # training data synthesis
+            self.gt = data['gt'].to(self.device)
+            self.gt_usm = self.usm_sharpener(self.gt)
+
+            self.kernel1 = data['kernel1'].to(self.device)
+            self.kernel2 = data['kernel2'].to(self.device)
+            self.sinc_kernel = data['sinc_kernel'].to(self.device)
+
+            ori_h, ori_w = self.gt.size()[2:4]
+
+            # ----------------------- The first degradation process ----------------------- #
+            # blur
+            out = filter2D(self.gt_usm, self.kernel1)
+            # random resize
+            updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
+            if updown_type == 'up':
+                scale = np.random.uniform(1, self.opt['resize_range'][1])
+            elif updown_type == 'down':
+                scale = np.random.uniform(self.opt['resize_range'][0], 1)
+            else:
+                scale = 1
+            mode = random.choice(['area', 'bilinear', 'bicubic'])
+            out = F.interpolate(out, scale_factor=scale, mode=mode)
+            # add noise
+            gray_noise_prob = self.opt['gray_noise_prob']
+            if np.random.uniform() < self.opt['gaussian_noise_prob']:
+                out = random_add_gaussian_noise_pt(
+                    out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
+            else:
+                out = random_add_poisson_noise_pt(
+                    out,
+                    scale_range=self.opt['poisson_scale_range'],
+                    gray_prob=gray_noise_prob,
+                    clip=True,
+                    rounds=False)
+            # JPEG compression
+            jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
+            out = torch.clamp(out, 0, 1)  # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
+            out = self.jpeger(out, quality=jpeg_p)
+
+            # ----------------------- The second degradation process ----------------------- #
+            # blur
+            if np.random.uniform() < self.opt['second_blur_prob']:
+                out = filter2D(out, self.kernel2)
+            # random resize
+            updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
+            if updown_type == 'up':
+                scale = np.random.uniform(1, self.opt['resize_range2'][1])
+            elif updown_type == 'down':
+                scale = np.random.uniform(self.opt['resize_range2'][0], 1)
+            else:
+                scale = 1
+            mode = random.choice(['area', 'bilinear', 'bicubic'])
+            out = F.interpolate(
+                out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
+            # add noise
+            gray_noise_prob = self.opt['gray_noise_prob2']
+            if np.random.uniform() < self.opt['gaussian_noise_prob2']:
+                out = random_add_gaussian_noise_pt(
+                    out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
+            else:
+                out = random_add_poisson_noise_pt(
+                    out,
+                    scale_range=self.opt['poisson_scale_range2'],
+                    gray_prob=gray_noise_prob,
+                    clip=True,
+                    rounds=False)
+
+            # JPEG compression + the final sinc filter
+            # We also need to resize images to desired sizes. We group [resize back + sinc filter] together
+            # as one operation.
+            # We consider two orders:
+            #   1. [resize back + sinc filter] + JPEG compression
+            #   2. JPEG compression + [resize back + sinc filter]
+            # Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
+            if np.random.uniform() < 0.5:
+                # resize back + the final sinc filter
+                mode = random.choice(['area', 'bilinear', 'bicubic'])
+                out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
+                out = filter2D(out, self.sinc_kernel)
+                # JPEG compression
+                jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
+                out = torch.clamp(out, 0, 1)
+                out = self.jpeger(out, quality=jpeg_p)
+            else:
+                # JPEG compression
+                jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
+                out = torch.clamp(out, 0, 1)
+                out = self.jpeger(out, quality=jpeg_p)
+                # resize back + the final sinc filter
+                mode = random.choice(['area', 'bilinear', 'bicubic'])
+                out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
+                out = filter2D(out, self.sinc_kernel)
+
+            # clamp and round
+            self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+
+            # random crop
+            gt_size = self.opt['gt_size']
+            (self.gt, self.gt_usm), self.lq = paired_random_crop([self.gt, self.gt_usm], self.lq, gt_size,
+                                                                 self.opt['scale'])
+
+            # training pair pool
+            self._dequeue_and_enqueue()
+            # sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
+            self.gt_usm = self.usm_sharpener(self.gt)
+            self.lq = self.lq.contiguous()  # for the warning: grad and param do not obey the gradient layout contract
+        else:
+            # for paired training or validation
+            self.lq = data['lq'].to(self.device)
+            if 'gt' in data:
+                self.gt = data['gt'].to(self.device)
+                self.gt_usm = self.usm_sharpener(self.gt)
+
+    def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
+        # do not use the synthetic process during validation
+        self.is_train = False
+        super(RealESRGANModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
+        self.is_train = True
+
+    def optimize_parameters(self, current_iter):
+        # usm sharpening
+        l1_gt = self.gt_usm
+        percep_gt = self.gt_usm
+        gan_gt = self.gt_usm
+        if self.opt['l1_gt_usm'] is False:
+            l1_gt = self.gt
+        if self.opt['percep_gt_usm'] is False:
+            percep_gt = self.gt
+        if self.opt['gan_gt_usm'] is False:
+            gan_gt = self.gt
+
+        # optimize net_g
+        for p in self.net_d.parameters():
+            p.requires_grad = False
+
+        self.optimizer_g.zero_grad()
+        self.output = self.net_g(self.lq)
+
+        l_g_total = 0
+        loss_dict = OrderedDict()
+        if (current_iter % self.net_d_iters == 0 and current_iter > self.net_d_init_iters):
+            # pixel loss
+            if self.cri_pix:
+                l_g_pix = self.cri_pix(self.output, l1_gt)
+                l_g_total += l_g_pix
+                loss_dict['l_g_pix'] = l_g_pix
+            # perceptual loss
+            if self.cri_perceptual:
+                l_g_percep, l_g_style = self.cri_perceptual(self.output, percep_gt)
+                if l_g_percep is not None:
+                    l_g_total += l_g_percep
+                    loss_dict['l_g_percep'] = l_g_percep
+                if l_g_style is not None:
+                    l_g_total += l_g_style
+                    loss_dict['l_g_style'] = l_g_style
+            # gan loss
+            fake_g_pred = self.net_d(self.output)
+            l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
+            l_g_total += l_g_gan
+            loss_dict['l_g_gan'] = l_g_gan
+
+            l_g_total.backward()
+            self.optimizer_g.step()
+
+        # optimize net_d
+        for p in self.net_d.parameters():
+            p.requires_grad = True
+
+        self.optimizer_d.zero_grad()
+        # real
+        real_d_pred = self.net_d(gan_gt)
+        l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
+        loss_dict['l_d_real'] = l_d_real
+        loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
+        l_d_real.backward()
+        # fake
+        fake_d_pred = self.net_d(self.output.detach().clone())  # clone for pt1.9
+        l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
+        loss_dict['l_d_fake'] = l_d_fake
+        loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
+        l_d_fake.backward()
+        self.optimizer_d.step()
+
+        if self.ema_decay > 0:
+            self.model_ema(decay=self.ema_decay)
+
+        self.log_dict = self.reduce_loss_dict(loss_dict)

JarvisIR/package/agent_tools/ESRGAN/realesrgan/models/realesrnet_model.py

@@ -0,0 +1,188 @@
+import numpy as np
+import random
+import torch
+from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
+from basicsr.data.transforms import paired_random_crop
+from basicsr.models.sr_model import SRModel
+from basicsr.utils import DiffJPEG, USMSharp
+from basicsr.utils.img_process_util import filter2D
+from basicsr.utils.registry import MODEL_REGISTRY
+from torch.nn import functional as F
+
+
+@MODEL_REGISTRY.register()
+class RealESRNetModel(SRModel):
+    """RealESRNet Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    It is trained without GAN losses.
+    It mainly performs:
+    1. randomly synthesize LQ images in GPU tensors
+    2. optimize the networks with GAN training.
+    """
+
+    def __init__(self, opt):
+        super(RealESRNetModel, self).__init__(opt)
+        self.jpeger = DiffJPEG(differentiable=False).cuda()  # simulate JPEG compression artifacts
+        self.usm_sharpener = USMSharp().cuda()  # do usm sharpening
+        self.queue_size = opt.get('queue_size', 180)
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self):
+        """It is the training pair pool for increasing the diversity in a batch.
+
+        Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
+        batch could not have different resize scaling factors. Therefore, we employ this training pair pool
+        to increase the degradation diversity in a batch.
+        """
+        # initialize
+        b, c, h, w = self.lq.size()
+        if not hasattr(self, 'queue_lr'):
+            assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
+            self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
+            _, c, h, w = self.gt.size()
+            self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
+            self.queue_ptr = 0
+        if self.queue_ptr == self.queue_size:  # the pool is full
+            # do dequeue and enqueue
+            # shuffle
+            idx = torch.randperm(self.queue_size)
+            self.queue_lr = self.queue_lr[idx]
+            self.queue_gt = self.queue_gt[idx]
+            # get first b samples
+            lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
+            gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
+            # update the queue
+            self.queue_lr[0:b, :, :, :] = self.lq.clone()
+            self.queue_gt[0:b, :, :, :] = self.gt.clone()
+
+            self.lq = lq_dequeue
+            self.gt = gt_dequeue
+        else:
+            # only do enqueue
+            self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
+            self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
+            self.queue_ptr = self.queue_ptr + b
+
+    @torch.no_grad()
+    def feed_data(self, data):
+        """Accept data from dataloader, and then add two-order degradations to obtain LQ images.
+        """
+        if self.is_train and self.opt.get('high_order_degradation', True):
+            # training data synthesis
+            self.gt = data['gt'].to(self.device)
+            # USM sharpen the GT images
+            if self.opt['gt_usm'] is True:
+                self.gt = self.usm_sharpener(self.gt)
+
+            self.kernel1 = data['kernel1'].to(self.device)
+            self.kernel2 = data['kernel2'].to(self.device)
+            self.sinc_kernel = data['sinc_kernel'].to(self.device)
+
+            ori_h, ori_w = self.gt.size()[2:4]
+
+            # ----------------------- The first degradation process ----------------------- #
+            # blur
+            out = filter2D(self.gt, self.kernel1)
+            # random resize
+            updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
+            if updown_type == 'up':
+                scale = np.random.uniform(1, self.opt['resize_range'][1])
+            elif updown_type == 'down':
+                scale = np.random.uniform(self.opt['resize_range'][0], 1)
+            else:
+                scale = 1
+            mode = random.choice(['area', 'bilinear', 'bicubic'])
+            out = F.interpolate(out, scale_factor=scale, mode=mode)
+            # add noise
+            gray_noise_prob = self.opt['gray_noise_prob']
+            if np.random.uniform() < self.opt['gaussian_noise_prob']:
+                out = random_add_gaussian_noise_pt(
+                    out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
+            else:
+                out = random_add_poisson_noise_pt(
+                    out,
+                    scale_range=self.opt['poisson_scale_range'],
+                    gray_prob=gray_noise_prob,
+                    clip=True,
+                    rounds=False)
+            # JPEG compression
+            jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
+            out = torch.clamp(out, 0, 1)  # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
+            out = self.jpeger(out, quality=jpeg_p)
+
+            # ----------------------- The second degradation process ----------------------- #
+            # blur
+            if np.random.uniform() < self.opt['second_blur_prob']:
+                out = filter2D(out, self.kernel2)
+            # random resize
+            updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
+            if updown_type == 'up':
+                scale = np.random.uniform(1, self.opt['resize_range2'][1])
+            elif updown_type == 'down':
+                scale = np.random.uniform(self.opt['resize_range2'][0], 1)
+            else:
+                scale = 1
+            mode = random.choice(['area', 'bilinear', 'bicubic'])
+            out = F.interpolate(
+                out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
+            # add noise
+            gray_noise_prob = self.opt['gray_noise_prob2']
+            if np.random.uniform() < self.opt['gaussian_noise_prob2']:
+                out = random_add_gaussian_noise_pt(
+                    out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
+            else:
+                out = random_add_poisson_noise_pt(
+                    out,
+                    scale_range=self.opt['poisson_scale_range2'],
+                    gray_prob=gray_noise_prob,
+                    clip=True,
+                    rounds=False)
+
+            # JPEG compression + the final sinc filter
+            # We also need to resize images to desired sizes. We group [resize back + sinc filter] together
+            # as one operation.
+            # We consider two orders:
+            #   1. [resize back + sinc filter] + JPEG compression
+            #   2. JPEG compression + [resize back + sinc filter]
+            # Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
+            if np.random.uniform() < 0.5:
+                # resize back + the final sinc filter
+                mode = random.choice(['area', 'bilinear', 'bicubic'])
+                out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
+                out = filter2D(out, self.sinc_kernel)
+                # JPEG compression
+                jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
+                out = torch.clamp(out, 0, 1)
+                out = self.jpeger(out, quality=jpeg_p)
+            else:
+                # JPEG compression
+                jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
+                out = torch.clamp(out, 0, 1)
+                out = self.jpeger(out, quality=jpeg_p)
+                # resize back + the final sinc filter
+                mode = random.choice(['area', 'bilinear', 'bicubic'])
+                out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
+                out = filter2D(out, self.sinc_kernel)
+
+            # clamp and round
+            self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+
+            # random crop
+            gt_size = self.opt['gt_size']
+            self.gt, self.lq = paired_random_crop(self.gt, self.lq, gt_size, self.opt['scale'])
+
+            # training pair pool
+            self._dequeue_and_enqueue()
+            self.lq = self.lq.contiguous()  # for the warning: grad and param do not obey the gradient layout contract
+        else:
+            # for paired training or validation
+            self.lq = data['lq'].to(self.device)
+            if 'gt' in data:
+                self.gt = data['gt'].to(self.device)
+                self.gt_usm = self.usm_sharpener(self.gt)
+
+    def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
+        # do not use the synthetic process during validation
+        self.is_train = False
+        super(RealESRNetModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
+        self.is_train = True

JarvisIR/package/agent_tools/ESRGAN/realesrgan/train.py

@@ -0,0 +1,11 @@
+# flake8: noqa
+import os.path as osp
+from basicsr.train import train_pipeline
+
+import realesrgan.archs
+import realesrgan.data
+import realesrgan.models
+
+if __name__ == '__main__':
+    root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
+    train_pipeline(root_path)

JarvisIR/package/agent_tools/ESRGAN/realesrgan/utils.py

@@ -0,0 +1,309 @@
+import cv2
+import math
+import numpy as np
+import os
+import queue
+import threading
+import torch
+from basicsr.utils.download_util import load_file_from_url
+from torch.nn import functional as F
+
+ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+class RealESRGANer():
+    """A helper class for upsampling images with RealESRGAN.
+
+    Args:
+        scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.
+        model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).
+        model (nn.Module): The defined network. Default: None.
+        tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop
+            input images into tiles, and then process each of them. Finally, they will be merged into one image.
+            0 denotes for do not use tile. Default: 0.
+        tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.
+        pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.
+        half (float): Whether to use half precision during inference. Default: False.
+    """
+
+    def __init__(self,
+                 scale,
+                 model_path,
+                 dni_weight=None,
+                 model=None,
+                 tile=0,
+                 tile_pad=10,
+                 pre_pad=10,
+                 half=False,
+                 device=None,
+                 gpu_id=None):
+        self.scale = scale
+        self.tile_size = tile
+        self.tile_pad = tile_pad
+        self.pre_pad = pre_pad
+        self.mod_scale = None
+        self.half = half
+
+        # initialize model
+        self.device = device
+
+        if isinstance(model_path, list):
+            # dni
+            assert len(model_path) == len(dni_weight), 'model_path and dni_weight should have the save length.'
+            loadnet = self.dni(model_path[0], model_path[1], dni_weight)
+        else:
+            # if the model_path starts with https, it will first download models to the folder: weights
+            if model_path.startswith('https://'):
+                model_path = load_file_from_url(
+                    url=model_path, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+            loadnet = torch.load(model_path, map_location=torch.device('cpu'))
+
+        # prefer to use params_ema
+        if 'params_ema' in loadnet:
+            keyname = 'params_ema'
+        else:
+            keyname = 'params'
+        model.load_state_dict(loadnet[keyname], strict=True)
+
+        model.eval()
+        self.model = model.to(self.device)
+        if self.half:
+            self.model = self.model.half()
+
+    def dni(self, net_a, net_b, dni_weight, key='params', loc='cpu'):
+        """Deep network interpolation.
+
+        ``Paper: Deep Network Interpolation for Continuous Imagery Effect Transition``
+        """
+        net_a = torch.load(net_a, map_location=torch.device(loc))
+        net_b = torch.load(net_b, map_location=torch.device(loc))
+        for k, v_a in net_a[key].items():
+            net_a[key][k] = dni_weight[0] * v_a + dni_weight[1] * net_b[key][k]
+        return net_a
+
+    def pre_process(self, img):
+        """Pre-process, such as pre-pad and mod pad, so that the images can be divisible
+        """
+        img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
+        self.img = img.unsqueeze(0).to(self.device)
+        if self.half:
+            self.img = self.img.half()
+
+        # pre_pad
+        if self.pre_pad != 0:
+            self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
+        # mod pad for divisible borders
+        if self.scale == 2:
+            self.mod_scale = 2
+        elif self.scale == 1:
+            self.mod_scale = 4
+        if self.mod_scale is not None:
+            self.mod_pad_h, self.mod_pad_w = 0, 0
+            _, _, h, w = self.img.size()
+            if (h % self.mod_scale != 0):
+                self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
+            if (w % self.mod_scale != 0):
+                self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
+            self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
+
+    def process(self):
+        # model inference
+        self.output = self.model(self.img)
+
+    def tile_process(self):
+        """It will first crop input images to tiles, and then process each tile.
+        Finally, all the processed tiles are merged into one images.
+
+        Modified from: https://github.com/ata4/esrgan-launcher
+        """
+        batch, channel, height, width = self.img.shape
+        output_height = height * self.scale
+        output_width = width * self.scale
+        output_shape = (batch, channel, output_height, output_width)
+
+        # start with black image
+        self.output = self.img.new_zeros(output_shape)
+        tiles_x = math.ceil(width / self.tile_size)
+        tiles_y = math.ceil(height / self.tile_size)
+
+        # loop over all tiles
+        for y in range(tiles_y):
+            for x in range(tiles_x):
+                # extract tile from input image
+                ofs_x = x * self.tile_size
+                ofs_y = y * self.tile_size
+                # input tile area on total image
+                input_start_x = ofs_x
+                input_end_x = min(ofs_x + self.tile_size, width)
+                input_start_y = ofs_y
+                input_end_y = min(ofs_y + self.tile_size, height)
+
+                # input tile area on total image with padding
+                input_start_x_pad = max(input_start_x - self.tile_pad, 0)
+                input_end_x_pad = min(input_end_x + self.tile_pad, width)
+                input_start_y_pad = max(input_start_y - self.tile_pad, 0)
+                input_end_y_pad = min(input_end_y + self.tile_pad, height)
+
+                # input tile dimensions
+                input_tile_width = input_end_x - input_start_x
+                input_tile_height = input_end_y - input_start_y
+                tile_idx = y * tiles_x + x + 1
+                input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
+
+                # upscale tile
+                try:
+                    with torch.no_grad():
+                        output_tile = self.model(input_tile)
+                except RuntimeError as error:
+                    print('Error', error)
+                print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
+
+                # output tile area on total image
+                output_start_x = input_start_x * self.scale
+                output_end_x = input_end_x * self.scale
+                output_start_y = input_start_y * self.scale
+                output_end_y = input_end_y * self.scale
+
+                # output tile area without padding
+                output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
+                output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
+                output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
+                output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
+
+                # put tile into output image
+                self.output[:, :, output_start_y:output_end_y,
+                            output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
+                                                                       output_start_x_tile:output_end_x_tile]
+
+    def post_process(self):
+        # remove extra pad
+        if self.mod_scale is not None:
+            _, _, h, w = self.output.size()
+            self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
+        # remove prepad
+        if self.pre_pad != 0:
+            _, _, h, w = self.output.size()
+            self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
+        return self.output
+
+    @torch.no_grad()
+    def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
+        h_input, w_input = img.shape[0:2]
+        # img: numpy
+        img = img.astype(np.float32)
+        if np.max(img) > 256:  # 16-bit image
+            max_range = 65535
+            print('\tInput is a 16-bit image')
+        else:
+            max_range = 255
+        img = img / max_range
+        if len(img.shape) == 2:  # gray image
+            img_mode = 'L'
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+        elif img.shape[2] == 4:  # RGBA image with alpha channel
+            img_mode = 'RGBA'
+            alpha = img[:, :, 3]
+            img = img[:, :, 0:3]
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            if alpha_upsampler == 'realesrgan':
+                alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
+        else:
+            img_mode = 'RGB'
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        # ------------------- process image (without the alpha channel) ------------------- #
+        self.pre_process(img)
+        if self.tile_size > 0:
+            self.tile_process()
+        else:
+            self.process()
+        output_img = self.post_process()
+        output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+        output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
+        if img_mode == 'L':
+            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
+
+        # ------------------- process the alpha channel if necessary ------------------- #
+        if img_mode == 'RGBA':
+            if alpha_upsampler == 'realesrgan':
+                self.pre_process(alpha)
+                if self.tile_size > 0:
+                    self.tile_process()
+                else:
+                    self.process()
+                output_alpha = self.post_process()
+                output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+                output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
+                output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
+            else:  # use the cv2 resize for alpha channel
+                h, w = alpha.shape[0:2]
+                output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
+
+            # merge the alpha channel
+            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
+            output_img[:, :, 3] = output_alpha
+
+        # ------------------------------ return ------------------------------ #
+        if max_range == 65535:  # 16-bit image
+            output = (output_img * 65535.0).round().astype(np.uint16)
+        else:
+            output = (output_img * 255.0).round().astype(np.uint8)
+
+        if outscale is not None and outscale != float(self.scale):
+            output = cv2.resize(
+                output, (
+                    int(w_input * outscale),
+                    int(h_input * outscale),
+                ), interpolation=cv2.INTER_LANCZOS4)
+
+        return output, img_mode
+
+
+class PrefetchReader(threading.Thread):
+    """Prefetch images.
+
+    Args:
+        img_list (list[str]): A image list of image paths to be read.
+        num_prefetch_queue (int): Number of prefetch queue.
+    """
+
+    def __init__(self, img_list, num_prefetch_queue):
+        super().__init__()
+        self.que = queue.Queue(num_prefetch_queue)
+        self.img_list = img_list
+
+    def run(self):
+        for img_path in self.img_list:
+            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+            self.que.put(img)
+
+        self.que.put(None)
+
+    def __next__(self):
+        next_item = self.que.get()
+        if next_item is None:
+            raise StopIteration
+        return next_item
+
+    def __iter__(self):
+        return self
+
+
+class IOConsumer(threading.Thread):
+
+    def __init__(self, opt, que, qid):
+        super().__init__()
+        self._queue = que
+        self.qid = qid
+        self.opt = opt
+
+    def run(self):
+        while True:
+            msg = self._queue.get()
+            if isinstance(msg, str) and msg == 'quit':
+                break
+
+            output = msg['output']
+            save_path = msg['save_path']
+            cv2.imwrite(save_path, output)
+        print(f'IO worker {self.qid} is done.')

JarvisIR/package/agent_tools/HVICIDNet/inference.py

@@ -0,0 +1,65 @@
+import os
+import numpy as np
+import torch
+from PIL import Image
+from .net.CIDNet import CIDNet
+import torchvision.transforms as transforms
+import torch.nn.functional as F
+import platform
+import argparse
+
+def load_hvicidnet_model(path, device):
+    model = CIDNet().to(device)
+    model.load_state_dict(torch.load(path, map_location=lambda storage, loc: storage))
+    model.eval()
+    return model
+
+def hvicidnet_predict(model, input_img, output_dir, device):
+    # Load image if path is provided as string
+    if isinstance(input_img, str):
+        img_name = os.path.basename(input_img).split('.')[0]
+        input_img = Image.open(input_img)
+    else:
+        img_name = "output" # default name if input is not a path
+
+    torch.set_grad_enabled(False)
+    pil2tensor = transforms.Compose([transforms.ToTensor()])
+    input = pil2tensor(input_img)
+    factor = 8
+    h, w = input.shape[1], input.shape[2]
+    H, W = ((h + factor) // factor) * factor, ((w + factor) // factor) * factor
+    padh = H - h if h % factor != 0 else 0
+    padw = W - w if w % factor != 0 else 0
+    input = F.pad(input.unsqueeze(0), (0,padw,0,padh), 'reflect')
+    gamma = 1
+    alpha_s = 1.0
+    alpha_i = 1.0
+    with torch.no_grad():
+        model.trans.alpha_s = alpha_s
+        model.trans.alpha = alpha_i
+        output = model(input.cuda()**gamma)
+
+    output = torch.clamp(output.to(device),0,1).cpu()
+    output = output[:, :, :h, :w]
+    enhanced_img = transforms.ToPILImage()(output.squeeze(0))
+    if isinstance(input_img, str):
+        original_img = Image.open(input_img)
+        enhanced_img = enhanced_img.resize(original_img.size, Image.LANCZOS)
+
+    # Save the output
+    os.makedirs(output_dir, exist_ok=True)
+    save_path = os.path.join(output_dir, img_name+'.png')
+    enhanced_img.save(save_path)
+    return save_path
+
+
+if __name__ == '__main__':
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model_path = "checkpoints/HVICIDNet/generalization.pth"
+    img_path = "./Test_Input/108.png"
+    output_folder = "./output"
+    model = load_hvicidnet_model(model_path, device)
+    save_path = hvicidnet_predict(model, img_path, output_folder, device)
+    print(f"processed image saved to: {save_path}")
+    
+

JarvisIR/package/agent_tools/HVICIDNet/loss/loss_utils.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn.functional as F
+import functools
+from math import exp
+from torch.autograd import Variable
+
+
+
+
+def reduce_loss(loss, reduction):
+    """Reduce loss as specified.
+
+    Args:
+        loss (Tensor): Elementwise loss tensor.
+        reduction (str): Options are 'none', 'mean' and 'sum'.
+
+    Returns:
+        Tensor: Reduced loss tensor.
+    """
+    reduction_enum = F._Reduction.get_enum(reduction)
+    # none: 0, elementwise_mean:1, sum: 2
+    if reduction_enum == 0:
+        return loss
+    elif reduction_enum == 1:
+        return loss.mean()
+    else:
+        return loss.sum()
+
+
+def weight_reduce_loss(loss, weight=None, reduction='mean'):
+    """Apply element-wise weight and reduce loss.
+
+    Args:
+        loss (Tensor): Element-wise loss.
+        weight (Tensor): Element-wise weights. Default: None.
+        reduction (str): Same as built-in losses of PyTorch. Options are
+            'none', 'mean' and 'sum'. Default: 'mean'.
+
+    Returns:
+        Tensor: Loss values.
+    """
+    # if weight is specified, apply element-wise weight
+    if weight is not None:
+        assert weight.dim() == loss.dim()
+        assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
+        loss = loss * weight
+
+    # if weight is not specified or reduction is sum, just reduce the loss
+    if weight is None or reduction == 'sum':
+        loss = reduce_loss(loss, reduction)
+    # if reduction is mean, then compute mean over weight region
+    elif reduction == 'mean':
+        if weight.size(1) > 1:
+            weight = weight.sum()
+        else:
+            weight = weight.sum() * loss.size(1)
+        loss = loss.sum() / weight
+
+    return loss
+
+def weighted_loss(loss_func):
+    """Create a weighted version of a given loss function.
+
+    To use this decorator, the loss function must have the signature like
+    `loss_func(pred, target, **kwargs)`. The function only needs to compute
+    element-wise loss without any reduction. This decorator will add weight
+    and reduction arguments to the function. The decorated function will have
+    the signature like `loss_func(pred, target, weight=None, reduction='mean',
+    **kwargs)`.
+
+    :Example:
+
+    >>> import torch
+    >>> @weighted_loss
+    >>> def l1_loss(pred, target):
+    >>>     return (pred - target).abs()
+
+    >>> pred = torch.Tensor([0, 2, 3])
+    >>> target = torch.Tensor([1, 1, 1])
+    >>> weight = torch.Tensor([1, 0, 1])
+
+    >>> l1_loss(pred, target)
+    tensor(1.3333)
+    >>> l1_loss(pred, target, weight)
+    tensor(1.5000)
+    >>> l1_loss(pred, target, reduction='none')
+    tensor([1., 1., 2.])
+    >>> l1_loss(pred, target, weight, reduction='sum')
+    tensor(3.)
+    """
+
+    @functools.wraps(loss_func)
+    def wrapper(pred, target, weight=None, reduction='mean', **kwargs):
+        # get element-wise loss
+        loss = loss_func(pred, target, **kwargs)
+        loss = weight_reduce_loss(loss, weight, reduction)
+        return loss
+
+    return wrapper
+
+@weighted_loss
+def l1_loss(pred, target):
+    return F.l1_loss(pred, target, reduction='none')
+
+@weighted_loss
+def mse_loss(pred, target):
+    return F.mse_loss(pred, target, reduction='none')
+
+
+
+
+
+def gaussian(window_size,sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss/torch.sum(gauss)  
+
+
+def create_window(window_size,channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)  
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+
+def map_ssim(img1, img2, window, window_size, channel, size_average=True):
+    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
+
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)

JarvisIR/package/agent_tools/HVICIDNet/loss/losses.py

@@ -0,0 +1,193 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from loss.vgg_arch import VGGFeatureExtractor, Registry
+from loss.loss_utils import *
+
+
+_reduction_modes = ['none', 'mean', 'sum']
+
+class L1Loss(nn.Module):
+    """L1 (mean absolute error, MAE) loss.
+
+    Args:
+        loss_weight (float): Loss weight for L1 loss. Default: 1.0.
+        reduction (str): Specifies the reduction to apply to the output.
+            Supported choices are 'none' | 'mean' | 'sum'. Default: 'mean'.
+    """
+
+    def __init__(self, loss_weight=1.0, reduction='mean'):
+        super(L1Loss, self).__init__()
+        if reduction not in ['none', 'mean', 'sum']:
+            raise ValueError(f'Unsupported reduction mode: {reduction}. '
+                             f'Supported ones are: {_reduction_modes}')
+
+        self.loss_weight = loss_weight
+        self.reduction = reduction
+
+    def forward(self, pred, target, weight=None, **kwargs):
+        """
+        Args:
+            pred (Tensor): of shape (N, C, H, W). Predicted tensor.
+            target (Tensor): of shape (N, C, H, W). Ground truth tensor.
+            weight (Tensor, optional): of shape (N, C, H, W). Element-wise
+                weights. Default: None.
+        """
+        return self.loss_weight * l1_loss(
+            pred, target, weight, reduction=self.reduction)
+        
+        
+        
+class EdgeLoss(nn.Module):
+    def __init__(self,loss_weight=1.0, reduction='mean'):
+        super(EdgeLoss, self).__init__()
+        k = torch.Tensor([[.05, .25, .4, .25, .05]])
+        self.kernel = torch.matmul(k.t(),k).unsqueeze(0).repeat(3,1,1,1).cuda()
+
+        self.weight = loss_weight
+        
+    def conv_gauss(self, img):
+        n_channels, _, kw, kh = self.kernel.shape
+        img = F.pad(img, (kw//2, kh//2, kw//2, kh//2), mode='replicate')
+        return F.conv2d(img, self.kernel, groups=n_channels)
+
+    def laplacian_kernel(self, current):
+        filtered    = self.conv_gauss(current)
+        down        = filtered[:,:,::2,::2]
+        new_filter  = torch.zeros_like(filtered)
+        new_filter[:,:,::2,::2] = down*4
+        filtered    = self.conv_gauss(new_filter)
+        diff = current - filtered
+        return diff
+
+    def forward(self, x, y):
+        loss = mse_loss(self.laplacian_kernel(x), self.laplacian_kernel(y))
+        return loss*self.weight
+
+
+class PerceptualLoss(nn.Module):
+    """Perceptual loss with commonly used style loss.
+
+    Args:
+        layer_weights (dict): The weight for each layer of vgg feature.
+            Here is an example: {'conv5_4': 1.}, which means the conv5_4
+            feature layer (before relu5_4) will be extracted with weight
+            1.0 in calculting losses.
+        vgg_type (str): The type of vgg network used as feature extractor.
+            Default: 'vgg19'.
+        use_input_norm (bool):  If True, normalize the input image in vgg.
+            Default: True.
+        range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
+            Default: False.
+        perceptual_weight (float): If `perceptual_weight > 0`, the perceptual
+            loss will be calculated and the loss will multiplied by the
+            weight. Default: 1.0.
+        style_weight (float): If `style_weight > 0`, the style loss will be
+            calculated and the loss will multiplied by the weight.
+            Default: 0.
+        criterion (str): Criterion used for perceptual loss. Default: 'l1'.
+    """
+
+    def __init__(self,
+                 layer_weights,
+                 vgg_type='vgg19',
+                 use_input_norm=True,
+                 range_norm=True,
+                 perceptual_weight=1.0,
+                 style_weight=0.,
+                 criterion='l1'):
+        super(PerceptualLoss, self).__init__()
+        self.perceptual_weight = perceptual_weight
+        self.style_weight = style_weight
+        self.layer_weights = layer_weights
+        self.vgg = VGGFeatureExtractor(
+            layer_name_list=list(layer_weights.keys()),
+            vgg_type=vgg_type,
+            use_input_norm=use_input_norm,
+            range_norm=range_norm)
+
+        self.criterion_type = criterion
+        if self.criterion_type == 'l1':
+            self.criterion = torch.nn.L1Loss()
+        elif self.criterion_type == 'l2':
+            self.criterion = torch.nn.L2loss()
+        elif self.criterion_type == 'mse':
+            self.criterion = torch.nn.MSELoss(reduction='mean')
+        elif self.criterion_type == 'fro':
+            self.criterion = None
+        else:
+            raise NotImplementedError(f'{criterion} criterion has not been supported.')
+
+    def forward(self, x, gt):
+        """Forward function.
+
+        Args:
+            x (Tensor): Input tensor with shape (n, c, h, w).
+            gt (Tensor): Ground-truth tensor with shape (n, c, h, w).
+
+        Returns:
+            Tensor: Forward results.
+        """
+        # extract vgg features
+        x_features = self.vgg(x)
+        gt_features = self.vgg(gt.detach())
+
+        # calculate perceptual loss
+        if self.perceptual_weight > 0:
+            percep_loss = 0
+            for k in x_features.keys():
+                if self.criterion_type == 'fro':
+                    percep_loss += torch.norm(x_features[k] - gt_features[k], p='fro') * self.layer_weights[k]
+                else:
+                    percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
+            percep_loss *= self.perceptual_weight
+        else:
+            percep_loss = None
+
+        # calculate style loss
+        if self.style_weight > 0:
+            style_loss = 0
+            for k in x_features.keys():
+                if self.criterion_type == 'fro':
+                    style_loss += torch.norm(
+                        self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), p='fro') * self.layer_weights[k]
+                else:
+                    style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
+                        gt_features[k])) * self.layer_weights[k]
+            style_loss *= self.style_weight
+        else:
+            style_loss = None
+
+        return percep_loss, style_loss
+
+
+
+
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True,weight=1.):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = 1
+        self.window = create_window(window_size, self.channel)
+        self.weight = weight
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.data.type() == img1.data.type():
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel)
+
+            if img1.is_cuda:
+                window = window.cuda(img1.get_device())
+            window = window.type_as(img1)
+
+            self.window = window
+            self.channel = channel
+
+        return (1. - map_ssim(img1, img2, window, self.window_size, channel, self.size_average)) * self.weight
+
+
+

JarvisIR/package/agent_tools/HVICIDNet/loss/niqe_pris_params.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a7c182a68c9e7f1b2e2e5ec723279d6f65d912b6fcaf37eb2bf03d7367c4296
+size 11850

JarvisIR/package/agent_tools/HVICIDNet/loss/niqe_utils.py

@@ -0,0 +1,559 @@
+
+import cv2
+import math
+import numpy as np
+from scipy.ndimage import convolve
+from scipy.special import gamma
+import torch
+
+def cubic(x):
+    """cubic function used for calculate_weights_indices."""
+    absx = torch.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    return (1.5 * absx3 - 2.5 * absx2 + 1) * (
+        (absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * (((absx > 1) *
+                                                                                     (absx <= 2)).type_as(absx))
+
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+    """Calculate weights and indices, used for imresize function.
+    Args:
+        in_length (int): Input length.
+        out_length (int): Output length.
+        scale (float): Scale factor.
+        kernel_width (int): Kernel width.
+        antialisaing (bool): Whether to apply anti-aliasing when downsampling.
+    """
+
+    if (scale < 1) and antialiasing:
+        # Use a modified kernel (larger kernel width) to simultaneously
+        # interpolate and antialias
+        kernel_width = kernel_width / scale
+
+    # Output-space coordinates
+    x = torch.linspace(1, out_length, out_length)
+
+    # Input-space coordinates. Calculate the inverse mapping such that 0.5
+    # in output space maps to 0.5 in input space, and 0.5 + scale in output
+    # space maps to 1.5 in input space.
+    u = x / scale + 0.5 * (1 - 1 / scale)
+
+    # What is the left-most pixel that can be involved in the computation?
+    left = torch.floor(u - kernel_width / 2)
+
+    # What is the maximum number of pixels that can be involved in the
+    # computation?  Note: it's OK to use an extra pixel here; if the
+    # corresponding weights are all zero, it will be eliminated at the end
+    # of this function.
+    p = math.ceil(kernel_width) + 2
+
+    # The indices of the input pixels involved in computing the k-th output
+    # pixel are in row k of the indices matrix.
+    indices = left.view(out_length, 1).expand(out_length, p) + torch.linspace(0, p - 1, p).view(1, p).expand(
+        out_length, p)
+
+    # The weights used to compute the k-th output pixel are in row k of the
+    # weights matrix.
+    distance_to_center = u.view(out_length, 1).expand(out_length, p) - indices
+
+    # apply cubic kernel
+    if (scale < 1) and antialiasing:
+        weights = scale * cubic(distance_to_center * scale)
+    else:
+        weights = cubic(distance_to_center)
+
+    # Normalize the weights matrix so that each row sums to 1.
+    weights_sum = torch.sum(weights, 1).view(out_length, 1)
+    weights = weights / weights_sum.expand(out_length, p)
+
+    # If a column in weights is all zero, get rid of it. only consider the
+    # first and last column.
+    weights_zero_tmp = torch.sum((weights == 0), 0)
+    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 1, p - 2)
+        weights = weights.narrow(1, 1, p - 2)
+    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 0, p - 2)
+        weights = weights.narrow(1, 0, p - 2)
+    weights = weights.contiguous()
+    indices = indices.contiguous()
+    sym_len_s = -indices.min() + 1
+    sym_len_e = indices.max() - in_length
+    indices = indices + sym_len_s - 1
+    return weights, indices, int(sym_len_s), int(sym_len_e)
+
+def imresize(img, scale, antialiasing=True):
+    """imresize function same as MATLAB.
+    It now only supports bicubic.
+    The same scale applies for both height and width.
+    Args:
+        img (Tensor | Numpy array):
+            Tensor: Input image with shape (c, h, w), [0, 1] range.
+            Numpy: Input image with shape (h, w, c), [0, 1] range.
+        scale (float): Scale factor. The same scale applies for both height
+            and width.
+        antialisaing (bool): Whether to apply anti-aliasing when downsampling.
+            Default: True.
+    Returns:
+        Tensor: Output image with shape (c, h, w), [0, 1] range, w/o round.
+    """
+    squeeze_flag = False
+    if type(img).__module__ == np.__name__:  # numpy type
+        numpy_type = True
+        if img.ndim == 2:
+            img = img[:, :, None]
+            squeeze_flag = True
+        img = torch.from_numpy(img.transpose(2, 0, 1)).float()
+    else:
+        numpy_type = False
+        if img.ndim == 2:
+            img = img.unsqueeze(0)
+            squeeze_flag = True
+
+    in_c, in_h, in_w = img.size()
+    out_h, out_w = math.ceil(in_h * scale), math.ceil(in_w * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # get weights and indices
+    weights_h, indices_h, sym_len_hs, sym_len_he = calculate_weights_indices(in_h, out_h, scale, kernel, kernel_width,
+                                                                             antialiasing)
+    weights_w, indices_w, sym_len_ws, sym_len_we = calculate_weights_indices(in_w, out_w, scale, kernel, kernel_width,
+                                                                             antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_c, in_h + sym_len_hs + sym_len_he, in_w)
+    img_aug.narrow(1, sym_len_hs, in_h).copy_(img)
+
+    sym_patch = img[:, :sym_len_hs, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, 0, sym_len_hs).copy_(sym_patch_inv)
+
+    sym_patch = img[:, -sym_len_he:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, sym_len_hs + in_h, sym_len_he).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(in_c, out_h, in_w)
+    kernel_width = weights_h.size(1)
+    for i in range(out_h):
+        idx = int(indices_h[i][0])
+        for j in range(in_c):
+            out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_h[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(in_c, out_h, in_w + sym_len_ws + sym_len_we)
+    out_1_aug.narrow(2, sym_len_ws, in_w).copy_(out_1)
+
+    sym_patch = out_1[:, :, :sym_len_ws]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, 0, sym_len_ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, :, -sym_len_we:]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, sym_len_ws + in_w, sym_len_we).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(in_c, out_h, out_w)
+    kernel_width = weights_w.size(1)
+    for i in range(out_w):
+        idx = int(indices_w[i][0])
+        for j in range(in_c):
+            out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_w[i])
+
+    if squeeze_flag:
+        out_2 = out_2.squeeze(0)
+    if numpy_type:
+        out_2 = out_2.numpy()
+        if not squeeze_flag:
+            out_2 = out_2.transpose(1, 2, 0)
+
+    return out_2
+
+
+def _convert_input_type_range(img):
+    """Convert the type and range of the input image.
+    It converts the input image to np.float32 type and range of [0, 1].
+    It is mainly used for pre-processing the input image in colorspace
+    conversion functions such as rgb2ycbcr and ycbcr2rgb.
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+    Returns:
+        (ndarray): The converted image with type of np.float32 and range of
+            [0, 1].
+    """
+    img_type = img.dtype
+    img = img.astype(np.float32)
+    if img_type == np.float32:
+        pass
+    elif img_type == np.uint8:
+        img /= 255.
+    else:
+        raise TypeError(f'The img type should be np.float32 or np.uint8, but got {img_type}')
+    return img
+
+
+def _convert_output_type_range(img, dst_type):
+    """Convert the type and range of the image according to dst_type.
+    It converts the image to desired type and range. If `dst_type` is np.uint8,
+    images will be converted to np.uint8 type with range [0, 255]. If
+    `dst_type` is np.float32, it converts the image to np.float32 type with
+    range [0, 1].
+    It is mainly used for post-processing images in colorspace conversion
+    functions such as rgb2ycbcr and ycbcr2rgb.
+    Args:
+        img (ndarray): The image to be converted with np.float32 type and
+            range [0, 255].
+        dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
+            converts the image to np.uint8 type with range [0, 255]. If
+            dst_type is np.float32, it converts the image to np.float32 type
+            with range [0, 1].
+    Returns:
+        (ndarray): The converted image with desired type and range.
+    """
+    if dst_type not in (np.uint8, np.float32):
+        raise TypeError(f'The dst_type should be np.float32 or np.uint8, but got {dst_type}')
+    if dst_type == np.uint8:
+        img = img.round()
+    else:
+        img /= 255.
+    return img.astype(dst_type)
+
+
+
+def rgb2ycbcr(img, y_only=False):
+    """Convert a RGB image to YCbCr image.
+    This function produces the same results as Matlab's `rgb2ycbcr` function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+    It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def bgr2ycbcr(img, y_only=False):
+    """Convert a BGR image to YCbCr image.
+    The bgr version of rgb2ycbcr.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+    It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+def ycbcr2rgb(img):
+    """Convert a YCbCr image to RGB image.
+    This function produces the same results as Matlab's ycbcr2rgb function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> RGB`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+    Returns:
+        ndarray: The converted RGB image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+                              [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]  # noqa: E126
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def to_y_channel(img):
+    """Change to Y channel of YCbCr.
+    Args:
+        img (ndarray): Images with range [0, 255].
+    Returns:
+        (ndarray): Images with range [0, 255] (float type) without round.
+    """
+    img = img.astype(np.float32) / 255.
+    if img.ndim == 3 and img.shape[2] == 3:
+        img = bgr2ycbcr(img, y_only=True)
+        img = img[..., None]
+    return img * 255.
+
+
+def reorder_image(img, input_order='HWC'):
+    """Reorder images to 'HWC' order.
+    If the input_order is (h, w), return (h, w, 1);
+    If the input_order is (c, h, w), return (h, w, c);
+    If the input_order is (h, w, c), return as it is.
+    Args:
+        img (ndarray): Input image.
+        input_order (str): Whether the input order is 'HWC' or 'CHW'.
+            If the input image shape is (h, w), input_order will not have
+            effects. Default: 'HWC'.
+    Returns:
+        ndarray: reordered image.
+    """
+
+    if input_order not in ['HWC', 'CHW']:
+        raise ValueError(f"Wrong input_order {input_order}. Supported input_orders are 'HWC' and 'CHW'")
+    if len(img.shape) == 2:
+        img = img[..., None]
+    if input_order == 'CHW':
+        img = img.transpose(1, 2, 0)
+    return img
+
+def rgb2ycbcr_pt(img, y_only=False):
+    """Convert RGB images to YCbCr images (PyTorch version).
+    It implements the ITU-R BT.601 conversion for standard-definition television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+    Args:
+        img (Tensor): Images with shape (n, 3, h, w), the range [0, 1], float, RGB format.
+         y_only (bool): Whether to only return Y channel. Default: False.
+    Returns:
+        (Tensor): converted images with the shape (n, 3/1, h, w), the range [0, 1], float.
+    """
+    if y_only:
+        weight = torch.tensor([[65.481], [128.553], [24.966]]).to(img)
+        out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + 16.0
+    else:
+        weight = torch.tensor([[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]).to(img)
+        bias = torch.tensor([16, 128, 128]).view(1, 3, 1, 1).to(img)
+        out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + bias
+
+    out_img = out_img / 255.
+    return
+
+def tensor2img(tensor):
+    im = (255. * tensor).data.cpu().numpy()
+    # clamp
+    im[im > 255] = 255
+    im[im < 0] = 0
+    im = im.astype(np.uint8)
+    return im
+
+def img2tensor(img):
+    img = (img / 255.).astype('float32')
+    if img.ndim ==2:
+        img = np.expand_dims(np.expand_dims(img, axis = 0),axis=0)
+    else:
+        img = np.transpose(img, (2, 0, 1))  # C, H, W
+        img = np.expand_dims(img, axis=0)
+    img = np.ascontiguousarray(img, dtype=np.float32)
+    tensor = torch.from_numpy(img)
+    return tensor
+
+def estimate_aggd_param(block):
+    """Estimate AGGD (Asymmetric Generalized Gaussian Distribution) parameters.
+    Args:
+        block (ndarray): 2D Image block.
+    Returns:
+        tuple: alpha (float), beta_l (float) and beta_r (float) for the AGGD
+            distribution (Estimating the parames in Equation 7 in the paper).
+    """
+    block = block.flatten()
+    gam = np.arange(0.2, 10.001, 0.001)  # len = 9801
+    gam_reciprocal = np.reciprocal(gam)
+    r_gam = np.square(gamma(gam_reciprocal * 2)) / (gamma(gam_reciprocal) * gamma(gam_reciprocal * 3))
+
+    left_std = np.sqrt(np.mean(block[block < 0]**2))
+    right_std = np.sqrt(np.mean(block[block > 0]**2))
+    gammahat = left_std / right_std
+    rhat = (np.mean(np.abs(block)))**2 / np.mean(block**2)
+    rhatnorm = (rhat * (gammahat**3 + 1) * (gammahat + 1)) / ((gammahat**2 + 1)**2)
+    array_position = np.argmin((r_gam - rhatnorm)**2)
+
+    alpha = gam[array_position]
+    beta_l = left_std * np.sqrt(gamma(1 / alpha) / gamma(3 / alpha))
+    beta_r = right_std * np.sqrt(gamma(1 / alpha) / gamma(3 / alpha))
+    return (alpha, beta_l, beta_r)
+
+
+def compute_feature(block):
+    """Compute features.
+    Args:
+        block (ndarray): 2D Image block.
+    Returns:
+        list: Features with length of 18.
+    """
+    feat = []
+    alpha, beta_l, beta_r = estimate_aggd_param(block)
+    feat.extend([alpha, (beta_l + beta_r) / 2])
+
+    # distortions disturb the fairly regular structure of natural images.
+    # This deviation can be captured by analyzing the sample distribution of
+    # the products of pairs of adjacent coefficients computed along
+    # horizontal, vertical and diagonal orientations.
+    shifts = [[0, 1], [1, 0], [1, 1], [1, -1]]
+    for i in range(len(shifts)):
+        shifted_block = np.roll(block, shifts[i], axis=(0, 1))
+        alpha, beta_l, beta_r = estimate_aggd_param(block * shifted_block)
+        # Eq. 8
+        mean = (beta_r - beta_l) * (gamma(2 / alpha) / gamma(1 / alpha))
+        feat.extend([alpha, mean, beta_l, beta_r])
+    return feat
+
+
+def niqe(img, mu_pris_param, cov_pris_param, gaussian_window, block_size_h=96, block_size_w=96):
+    """Calculate NIQE (Natural Image Quality Evaluator) metric.
+    ``Paper: Making a "Completely Blind" Image Quality Analyzer``
+    This implementation could produce almost the same results as the official
+    MATLAB codes: http://live.ece.utexas.edu/research/quality/niqe_release.zip
+    Note that we do not include block overlap height and width, since they are
+    always 0 in the official implementation.
+    For good performance, it is advisable by the official implementation to
+    divide the distorted image in to the same size patched as used for the
+    construction of multivariate Gaussian model.
+    Args:
+        img (ndarray): Input image whose quality needs to be computed. The
+            image must be a gray or Y (of YCbCr) image with shape (h, w).
+            Range [0, 255] with float type.
+        mu_pris_param (ndarray): Mean of a pre-defined multivariate Gaussian
+            model calculated on the pristine dataset.
+        cov_pris_param (ndarray): Covariance of a pre-defined multivariate
+            Gaussian model calculated on the pristine dataset.
+        gaussian_window (ndarray): A 7x7 Gaussian window used for smoothing the
+            image.
+        block_size_h (int): Height of the blocks in to which image is divided.
+            Default: 96 (the official recommended value).
+        block_size_w (int): Width of the blocks in to which image is divided.
+            Default: 96 (the official recommended value).
+    """
+    assert img.ndim == 2, ('Input image must be a gray or Y (of YCbCr) image with shape (h, w).')
+    # crop image
+    h, w = img.shape
+    num_block_h = math.floor(h / block_size_h)
+    num_block_w = math.floor(w / block_size_w)
+    img = img[0:num_block_h * block_size_h, 0:num_block_w * block_size_w]
+
+    distparam = []  # dist param is actually the multiscale features
+    for scale in (1, 2):  # perform on two scales (1, 2)
+        mu = convolve(img, gaussian_window, mode='nearest')
+        sigma = np.sqrt(np.abs(convolve(np.square(img), gaussian_window, mode='nearest') - np.square(mu)))
+        # normalize, as in Eq. 1 in the paper
+        img_nomalized = (img - mu) / (sigma + 1)
+
+        feat = []
+        for idx_w in range(num_block_w):
+            for idx_h in range(num_block_h):
+                # process ecah block
+                block = img_nomalized[idx_h * block_size_h // scale:(idx_h + 1) * block_size_h // scale,
+                                      idx_w * block_size_w // scale:(idx_w + 1) * block_size_w // scale]
+                feat.append(compute_feature(block))
+
+        distparam.append(np.array(feat))
+
+        if scale == 1:
+            img = imresize(img / 255., scale=0.5, antialiasing=True)
+            img = img * 255.
+
+    distparam = np.concatenate(distparam, axis=1)
+
+    # fit a MVG (multivariate Gaussian) model to distorted patch features
+    mu_distparam = np.nanmean(distparam, axis=0)
+    # use nancov. ref: https://ww2.mathworks.cn/help/stats/nancov.html
+    distparam_no_nan = distparam[~np.isnan(distparam).any(axis=1)]
+    cov_distparam = np.cov(distparam_no_nan, rowvar=False)
+
+    # compute niqe quality, Eq. 10 in the paper
+    invcov_param = np.linalg.pinv((cov_pris_param + cov_distparam) / 2)
+    quality = np.matmul(
+        np.matmul((mu_pris_param - mu_distparam), invcov_param), np.transpose((mu_pris_param - mu_distparam)))
+
+    quality = np.sqrt(quality)
+    quality = float(np.squeeze(quality))
+    return quality
+
+
+def calculate_niqe(img, crop_border=0,input_order='HWC', convert_to='y', **kwargs):
+    """Calculate NIQE (Natural Image Quality Evaluator) metric.
+    ``Paper: Making a "Completely Blind" Image Quality Analyzer``
+    This implementation could produce almost the same results as the official
+    MATLAB codes: http://live.ece.utexas.edu/research/quality/niqe_release.zip
+    > MATLAB R2021a result for tests/data/baboon.png: 5.72957338 (5.7296)
+    > Our re-implementation result for tests/data/baboon.png: 5.7295763 (5.7296)
+    We use the official params estimated from the pristine dataset.
+    We use the recommended block size (96, 96) without overlaps.
+    Args:
+        img (ndarray): Input image whose quality needs to be computed.
+            The input image must be in range [0, 255] with float/int type.
+            The input_order of image can be 'HW' or 'HWC' or 'CHW'. (BGR order)
+            If the input order is 'HWC' or 'CHW', it will be converted to gray
+            or Y (of YCbCr) image according to the ``convert_to`` argument.
+        crop_border (int): Cropped pixels in each edge of an image. These
+            pixels are not involved in the metric calculation.
+        input_order (str): Whether the input order is 'HW', 'HWC' or 'CHW'.
+            Default: 'HWC'.
+        convert_to (str): Whether converted to 'y' (of MATLAB YCbCr) or 'gray'.
+            Default: 'y'.
+    Returns:
+        float: NIQE result.
+    """
+    # ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+    # we use the official params estimated from the pristine dataset.
+    niqe_pris_params = np.load('./loss/niqe_pris_params.npz')
+    mu_pris_param = niqe_pris_params['mu_pris_param']
+    cov_pris_param = niqe_pris_params['cov_pris_param']
+    gaussian_window = niqe_pris_params['gaussian_window']
+
+    img = img.astype(np.float32)
+    if input_order != 'HW':
+        img = reorder_image(img, input_order=input_order)
+        if convert_to == 'y':
+            img = to_y_channel(img)
+        elif convert_to == 'gray':
+            img = cv2.cvtColor(img / 255., cv2.COLOR_BGR2GRAY) * 255.
+        img = np.squeeze(img)
+
+    if crop_border != 0:
+        img = img[crop_border:-crop_border, crop_border:-crop_border]
+
+    # round is necessary for being consistent with MATLAB's result
+    img = img.round()
+
+    niqe_result = niqe(img, mu_pris_param, cov_pris_param, gaussian_window)
+
+    return niqe_result

JarvisIR/package/agent_tools/HVICIDNet/loss/vgg_arch.py

@@ -0,0 +1,239 @@
+import os
+import torch
+from collections import OrderedDict
+from torch import nn as nn
+from torchvision.models import vgg as vgg
+
+class Registry():
+    """
+    The registry that provides name -> object mapping, to support third-party
+    users' custom modules.
+
+    To create a registry (e.g. a backbone registry):
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY = Registry('BACKBONE')
+
+    To register an object:
+
+    .. code-block:: python
+
+        @BACKBONE_REGISTRY.register()
+        class MyBackbone():
+            ...
+
+    Or:
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY.register(MyBackbone)
+    """
+
+    def __init__(self, name):
+        """
+        Args:
+            name (str): the name of this registry
+        """
+        self._name = name
+        self._obj_map = {}
+
+    def _do_register(self, name, obj):
+        assert (name not in self._obj_map), (f"An object named '{name}' was already registered "
+                                             f"in '{self._name}' registry!")
+        self._obj_map[name] = obj
+
+    def register(self, obj=None):
+        """
+        Register the given object under the the name `obj.__name__`.
+        Can be used as either a decorator or not.
+        See docstring of this class for usage.
+        """
+        if obj is None:
+            # used as a decorator
+            def deco(func_or_class):
+                name = func_or_class.__name__
+                self._do_register(name, func_or_class)
+                return func_or_class
+
+            return deco
+
+        # used as a function call
+        name = obj.__name__
+        self._do_register(name, obj)
+
+    def get(self, name):
+        ret = self._obj_map.get(name)
+        if ret is None:
+            raise KeyError(f"No object named '{name}' found in '{self._name}' registry!")
+        return ret
+
+    def __contains__(self, name):
+        return name in self._obj_map
+
+    def __iter__(self):
+        return iter(self._obj_map.items())
+
+    def keys(self):
+        return self._obj_map.keys()
+
+
+DATASET_REGISTRY = Registry('dataset')
+ARCH_REGISTRY = Registry('arch')
+MODEL_REGISTRY = Registry('model')
+LOSS_REGISTRY = Registry('loss')
+METRIC_REGISTRY = Registry('metric')
+
+VGG_PRETRAIN_PATH = 'experiments/pretrained_models/vgg19-dcbb9e9d.pth'
+NAMES = {
+    'vgg11': [
+        'conv1_1', 'relu1_1', 'pool1', 'conv2_1', 'relu2_1', 'pool2', 'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2',
+        'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2',
+        'pool5'
+    ],
+    'vgg13': [
+        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'pool4',
+        'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'pool5'
+    ],
+    'vgg16': [
+        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2',
+        'relu4_2', 'conv4_3', 'relu4_3', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3',
+        'pool5'
+    ],
+    'vgg19': [
+        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3', 'conv4_1',
+        'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4', 'conv5_1', 'relu5_1',
+        'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5'
+    ]
+}
+
+
+def insert_bn(names):
+    """Insert bn layer after each conv.
+
+    Args:
+        names (list): The list of layer names.
+
+    Returns:
+        list: The list of layer names with bn layers.
+    """
+    names_bn = []
+    for name in names:
+        names_bn.append(name)
+        if 'conv' in name:
+            position = name.replace('conv', '')
+            names_bn.append('bn' + position)
+    return names_bn
+
+
+@ARCH_REGISTRY.register()
+class VGGFeatureExtractor(nn.Module):
+    """VGG network for feature extraction.
+
+    In this implementation, we allow users to choose whether use normalization
+    in the input feature and the type of vgg network. Note that the pretrained
+    path must fit the vgg type.
+
+    Args:
+        layer_name_list (list[str]): Forward function returns the corresponding
+            features according to the layer_name_list.
+            Example: {'relu1_1', 'relu2_1', 'relu3_1'}.
+        vgg_type (str): Set the type of vgg network. Default: 'vgg19'.
+        use_input_norm (bool): If True, normalize the input image. Importantly,
+            the input feature must in the range [0, 1]. Default: True.
+        range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
+            Default: False.
+        requires_grad (bool): If true, the parameters of VGG network will be
+            optimized. Default: False.
+        remove_pooling (bool): If true, the max pooling operations in VGG net
+            will be removed. Default: False.
+        pooling_stride (int): The stride of max pooling operation. Default: 2.
+    """
+
+    def __init__(self,
+                 layer_name_list,
+                 vgg_type='vgg19',
+                 use_input_norm=True,
+                 range_norm=False,
+                 requires_grad=False,
+                 remove_pooling=False,
+                 pooling_stride=2):
+        super(VGGFeatureExtractor, self).__init__()
+
+        self.layer_name_list = layer_name_list
+        self.use_input_norm = use_input_norm
+        self.range_norm = range_norm
+
+        self.names = NAMES[vgg_type.replace('_bn', '')]
+        if 'bn' in vgg_type:
+            self.names = insert_bn(self.names)
+
+        # only borrow layers that will be used to avoid unused params
+        max_idx = 0
+        for v in layer_name_list:
+            idx = self.names.index(v)
+            if idx > max_idx:
+                max_idx = idx
+
+        if os.path.exists(VGG_PRETRAIN_PATH):
+            vgg_net = getattr(vgg, vgg_type)(pretrained=False)
+            state_dict = torch.load(VGG_PRETRAIN_PATH, map_location=lambda storage, loc: storage)
+            vgg_net.load_state_dict(state_dict)
+        else:
+            vgg_net = getattr(vgg, vgg_type)(pretrained=True)
+
+        features = vgg_net.features[:max_idx + 1]
+
+        modified_net = OrderedDict()
+        for k, v in zip(self.names, features):
+            if 'pool' in k:
+                # if remove_pooling is true, pooling operation will be removed
+                if remove_pooling:
+                    continue
+                else:
+                    # in some cases, we may want to change the default stride
+                    modified_net[k] = nn.MaxPool2d(kernel_size=2, stride=pooling_stride)
+            else:
+                modified_net[k] = v
+
+        self.vgg_net = nn.Sequential(modified_net).cuda()
+
+        if not requires_grad:
+            self.vgg_net.eval()
+            for param in self.parameters():
+                param.requires_grad = False
+        else:
+            self.vgg_net.train()
+            for param in self.parameters():
+                param.requires_grad = True
+
+        if self.use_input_norm:
+            # the mean is for image with range [0, 1]
+            self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda())
+            # the std is for image with range [0, 1]
+            self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda())
+
+    def forward(self, x):
+        """Forward function.
+
+        Args:
+            x (Tensor): Input tensor with shape (n, c, h, w).
+
+        Returns:
+            Tensor: Forward results.
+        """
+        if self.range_norm:
+            x = (x + 1) / 2
+        if self.use_input_norm:
+            x = (x - self.mean) / self.std
+        output = {}
+
+        for key, layer in self.vgg_net._modules.items():
+            x = layer(x)
+            if key in self.layer_name_list:
+                output[key] = x.clone()
+
+        return output

JarvisIR/package/agent_tools/HVICIDNet/mods.py

@@ -0,0 +1,153 @@
+import torch
+import torch.nn as nn
+import warnings
+import math
+
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+
+class cross_attention(nn.Module):
+    def __init__(self, dim, num_heads, dropout=0.):
+        super(cross_attention, self).__init__()
+        if dim % num_heads != 0:
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (dim, num_heads)
+            )
+        self.num_heads = num_heads
+        self.attention_head_size = int(dim / num_heads)
+
+        self.query = Depth_conv(in_ch=dim, out_ch=dim)
+        self.key = Depth_conv(in_ch=dim, out_ch=dim)
+        self.value = Depth_conv(in_ch=dim, out_ch=dim)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def transpose_for_scores(self, x):
+        '''
+        new_x_shape = x.size()[:-1] + (
+            self.num_heads,
+            self.attention_head_size,
+        )
+        print(new_x_shape)
+        x = x.view(*new_x_shape)
+        '''
+        return x.permute(0, 2, 1, 3)
+
+    def forward(self, hidden_states, ctx):
+        mixed_query_layer = self.query(hidden_states)
+        mixed_key_layer = self.key(ctx)
+        mixed_value_layer = self.value(ctx)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+        key_layer = self.transpose_for_scores(mixed_key_layer)
+        value_layer = self.transpose_for_scores(mixed_value_layer)
+
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        attention_probs = self.dropout(attention_probs)
+
+        ctx_layer = torch.matmul(attention_probs, value_layer)
+        ctx_layer = ctx_layer.permute(0, 2, 1, 3).contiguous()
+
+        return ctx_layer
+
+
+class Depth_conv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(Depth_conv, self).__init__()
+        self.depth_conv = nn.Conv2d(
+            in_channels=in_ch,
+            out_channels=in_ch,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=1,
+            groups=in_ch
+        )
+        self.point_conv = nn.Conv2d(
+            in_channels=in_ch,
+            out_channels=out_ch,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=0,
+            groups=1
+        )
+
+    def forward(self, input):
+        out = self.depth_conv(input)
+        out = self.point_conv(out)
+        return out
+
+
+class Dilated_Resblock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Dilated_Resblock, self).__init__()
+
+        sequence = list()
+        sequence += [
+            nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1),
+                      padding=1, dilation=(1, 1)),
+            nn.LeakyReLU(),
+            nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1),
+                      padding=2, dilation=(2, 2)),
+            nn.LeakyReLU(),
+            nn.Conv2d(out_channels, in_channels, kernel_size=(3, 3), stride=(1, 1),
+                      padding=3, dilation=(3, 3)),
+            nn.LeakyReLU(),
+            nn.Conv2d(out_channels, in_channels, kernel_size=(3, 3), stride=(1, 1),
+                      padding=2, dilation=(2, 2)),
+            nn.LeakyReLU(),
+            nn.Conv2d(out_channels, in_channels, kernel_size=(3, 3), stride=(1, 1),
+                      padding=1, dilation=(1, 1))
+
+        ]
+
+        self.model = nn.Sequential(*sequence)
+
+    def forward(self, x):
+        out = self.model(x) + x
+
+        return out
+
+
+class HFRM(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(HFRM, self).__init__()
+
+        self.conv_head = Depth_conv(in_channels, out_channels)
+
+        self.dilated_block_LH = Dilated_Resblock(out_channels, out_channels)
+        self.dilated_block_HL = Dilated_Resblock(out_channels, out_channels)
+
+        self.cross_attention0 = cross_attention(out_channels, num_heads=8)
+        self.dilated_block_HH = Dilated_Resblock(out_channels, out_channels)
+        self.conv_HH = nn.Conv2d(out_channels*2, out_channels, kernel_size=3, stride=1, padding=1)
+        self.cross_attention1 = cross_attention(out_channels, num_heads=8)
+
+        self.conv_tail = Depth_conv(out_channels, in_channels)
+
+    def forward(self, x):
+
+        b, c, h, w = x.shape
+
+        residual = x
+
+        x = self.conv_head(x)
+
+        x_HL, x_LH, x_HH = x[:b//3, ...], x[b//3:2*b//3, ...], x[2*b//3:, ...]
+
+        x_HH_LH = self.cross_attention0(x_LH, x_HH)
+        x_HH_HL = self.cross_attention1(x_HL, x_HH)
+
+        x_HL = self.dilated_block_HL(x_HL)
+        x_LH = self.dilated_block_LH(x_LH)
+
+        x_HH = self.dilated_block_HH(self.conv_HH(torch.cat((x_HH_LH, x_HH_HL), dim=1)))
+
+        out = self.conv_tail(torch.cat((x_HL, x_LH, x_HH), dim=0))
+
+        return out + residual

JarvisIR/package/agent_tools/HVICIDNet/net/CIDNet.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+from .HVI_transform import RGB_HVI
+from .transformer_utils import *
+from .LCA import *
+from huggingface_hub import PyTorchModelHubMixin
+
+class CIDNet(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, 
+                 channels=[36, 36, 72, 144],
+                 heads=[1, 2, 4, 8],
+                 norm=False
+        ):
+        super(CIDNet, self).__init__()
+        
+        
+        [ch1, ch2, ch3, ch4] = channels
+        [head1, head2, head3, head4] = heads
+        
+        # HV_ways
+        self.HVE_block0 = nn.Sequential(
+            nn.ReplicationPad2d(1),
+            nn.Conv2d(3, ch1, 3, stride=1, padding=0,bias=False)
+            )
+        self.HVE_block1 = NormDownsample(ch1, ch2, use_norm = norm)
+        self.HVE_block2 = NormDownsample(ch2, ch3, use_norm = norm)
+        self.HVE_block3 = NormDownsample(ch3, ch4, use_norm = norm)
+        
+        self.HVD_block3 = NormUpsample(ch4, ch3, use_norm = norm)
+        self.HVD_block2 = NormUpsample(ch3, ch2, use_norm = norm)
+        self.HVD_block1 = NormUpsample(ch2, ch1, use_norm = norm)
+        self.HVD_block0 = nn.Sequential(
+            nn.ReplicationPad2d(1),
+            nn.Conv2d(ch1, 2, 3, stride=1, padding=0,bias=False)
+        )
+        
+        
+        # I_ways
+        self.IE_block0 = nn.Sequential(
+            nn.ReplicationPad2d(1),
+            nn.Conv2d(1, ch1, 3, stride=1, padding=0,bias=False),
+            )
+        self.IE_block1 = NormDownsample(ch1, ch2, use_norm = norm)
+        self.IE_block2 = NormDownsample(ch2, ch3, use_norm = norm)
+        self.IE_block3 = NormDownsample(ch3, ch4, use_norm = norm)
+        
+        self.ID_block3 = NormUpsample(ch4, ch3, use_norm=norm)
+        self.ID_block2 = NormUpsample(ch3, ch2, use_norm=norm)
+        self.ID_block1 = NormUpsample(ch2, ch1, use_norm=norm)
+        self.ID_block0 =  nn.Sequential(
+            nn.ReplicationPad2d(1),
+            nn.Conv2d(ch1, 1, 3, stride=1, padding=0,bias=False),
+            )
+        
+        self.HV_LCA1 = HV_LCA(ch2, head2)
+        self.HV_LCA2 = HV_LCA(ch3, head3)
+        self.HV_LCA3 = HV_LCA(ch4, head4)
+        self.HV_LCA4 = HV_LCA(ch4, head4)
+        self.HV_LCA5 = HV_LCA(ch3, head3)
+        self.HV_LCA6 = HV_LCA(ch2, head2)
+        
+        self.I_LCA1 = I_LCA(ch2, head2)
+        self.I_LCA2 = I_LCA(ch3, head3)
+        self.I_LCA3 = I_LCA(ch4, head4)
+        self.I_LCA4 = I_LCA(ch4, head4)
+        self.I_LCA5 = I_LCA(ch3, head3)
+        self.I_LCA6 = I_LCA(ch2, head2)
+        
+        self.trans = RGB_HVI()
+        
+    def forward(self, x):
+        dtypes = x.dtype
+        hvi = self.trans.HVIT(x)
+        i = hvi[:,2,:,:].unsqueeze(1).to(dtypes)
+        # low
+        i_enc0 = self.IE_block0(i)
+        i_enc1 = self.IE_block1(i_enc0)
+        hv_0 = self.HVE_block0(hvi)
+        hv_1 = self.HVE_block1(hv_0)
+        i_jump0 = i_enc0
+        hv_jump0 = hv_0
+        
+        i_enc2 = self.I_LCA1(i_enc1, hv_1)
+        hv_2 = self.HV_LCA1(hv_1, i_enc1)
+        v_jump1 = i_enc2
+        hv_jump1 = hv_2
+        i_enc2 = self.IE_block2(i_enc2)
+        hv_2 = self.HVE_block2(hv_2)
+        
+        i_enc3 = self.I_LCA2(i_enc2, hv_2)
+        hv_3 = self.HV_LCA2(hv_2, i_enc2)
+        v_jump2 = i_enc3
+        hv_jump2 = hv_3
+        i_enc3 = self.IE_block3(i_enc2)
+        hv_3 = self.HVE_block3(hv_2)
+        
+        i_enc4 = self.I_LCA3(i_enc3, hv_3)
+        hv_4 = self.HV_LCA3(hv_3, i_enc3)
+        
+        i_dec4 = self.I_LCA4(i_enc4,hv_4)
+        hv_4 = self.HV_LCA4(hv_4, i_enc4)
+        
+        hv_3 = self.HVD_block3(hv_4, hv_jump2)
+        i_dec3 = self.ID_block3(i_dec4, v_jump2)
+        i_dec2 = self.I_LCA5(i_dec3, hv_3)
+        hv_2 = self.HV_LCA5(hv_3, i_dec3)
+        
+        hv_2 = self.HVD_block2(hv_2, hv_jump1)
+        i_dec2 = self.ID_block2(i_dec3, v_jump1)
+        
+        i_dec1 = self.I_LCA6(i_dec2, hv_2)
+        hv_1 = self.HV_LCA6(hv_2, i_dec2)
+        
+        i_dec1 = self.ID_block1(i_dec1, i_jump0)
+        i_dec0 = self.ID_block0(i_dec1)
+        hv_1 = self.HVD_block1(hv_1, hv_jump0)
+        hv_0 = self.HVD_block0(hv_1)
+        
+        output_hvi = torch.cat([hv_0, i_dec0], dim=1) + hvi
+        output_rgb = self.trans.PHVIT(output_hvi)
+
+        return output_rgb
+    
+    def HVIT(self,x):
+        hvi = self.trans.HVIT(x)
+        return hvi
+    
+    
+

JarvisIR/package/agent_tools/HVICIDNet/net/HVI_transform.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn as nn
+
+pi = 3.141592653589793
+
+class RGB_HVI(nn.Module):
+    def __init__(self):
+        super(RGB_HVI, self).__init__()
+        self.density_k = torch.nn.Parameter(torch.full([1],0.2)) # k is reciprocal to the paper mentioned
+        self.gated = False
+        self.gated2= False
+        self.alpha = 1.0
+        self.alpha_s = 1.3
+        self.this_k = 0
+        
+    def HVIT(self, img):
+        eps = 1e-8
+        device = img.device
+        dtypes = img.dtype
+        hue = torch.Tensor(img.shape[0], img.shape[2], img.shape[3]).to(device).to(dtypes)
+        value = img.max(1)[0].to(dtypes)
+        img_min = img.min(1)[0].to(dtypes)
+        hue[img[:,2]==value] = 4.0 + ( (img[:,0]-img[:,1]) / (value - img_min + eps)) [img[:,2]==value]
+        hue[img[:,1]==value] = 2.0 + ( (img[:,2]-img[:,0]) / (value - img_min + eps)) [img[:,1]==value]
+        hue[img[:,0]==value] = (0.0 + ((img[:,1]-img[:,2]) / (value - img_min + eps)) [img[:,0]==value]) % 6
+
+        hue[img.min(1)[0]==value] = 0.0
+        hue = hue/6.0
+
+        saturation = (value - img_min ) / (value + eps )
+        saturation[value==0] = 0
+
+        hue = hue.unsqueeze(1)
+        saturation = saturation.unsqueeze(1)
+        value = value.unsqueeze(1)
+        
+        k = self.density_k
+        self.this_k = k.item()
+        
+        color_sensitive = ((value * 0.5 * pi).sin() + eps).pow(k)
+        ch = (2.0 * pi * hue).cos()
+        cv = (2.0 * pi * hue).sin()
+        H = color_sensitive * saturation * ch
+        V = color_sensitive * saturation * cv
+        I = value
+        xyz = torch.cat([H, V, I],dim=1)
+        return xyz
+    
+    def PHVIT(self, img):
+        eps = 1e-8
+        H,V,I = img[:,0,:,:],img[:,1,:,:],img[:,2,:,:]
+        
+        # clip
+        H = torch.clamp(H,-1,1)
+        V = torch.clamp(V,-1,1)
+        I = torch.clamp(I,0,1)
+        
+        v = I
+        k = self.this_k
+        color_sensitive = ((v * 0.5 * pi).sin() + eps).pow(k)
+        H = (H) / (color_sensitive + eps)
+        V = (V) / (color_sensitive + eps)
+        H = torch.clamp(H,-1,1)
+        V = torch.clamp(V,-1,1)
+        h = torch.atan2(V + eps,H + eps) / (2*pi)
+        h = h%1
+        s = torch.sqrt(H**2 + V**2 + eps)
+        
+        if self.gated:
+            s = s * self.alpha_s
+        
+        s = torch.clamp(s,0,1)
+        v = torch.clamp(v,0,1)
+        
+        r = torch.zeros_like(h)
+        g = torch.zeros_like(h)
+        b = torch.zeros_like(h)
+        
+        hi = torch.floor(h * 6.0)
+        f = h * 6.0 - hi
+        p = v * (1. - s)
+        q = v * (1. - (f * s))
+        t = v * (1. - ((1. - f) * s))
+        
+        hi0 = hi==0
+        hi1 = hi==1
+        hi2 = hi==2
+        hi3 = hi==3
+        hi4 = hi==4
+        hi5 = hi==5
+        
+        r[hi0] = v[hi0]
+        g[hi0] = t[hi0]
+        b[hi0] = p[hi0]
+        
+        r[hi1] = q[hi1]
+        g[hi1] = v[hi1]
+        b[hi1] = p[hi1]
+        
+        r[hi2] = p[hi2]
+        g[hi2] = v[hi2]
+        b[hi2] = t[hi2]
+        
+        r[hi3] = p[hi3]
+        g[hi3] = q[hi3]
+        b[hi3] = v[hi3]
+        
+        r[hi4] = t[hi4]
+        g[hi4] = p[hi4]
+        b[hi4] = v[hi4]
+        
+        r[hi5] = v[hi5]
+        g[hi5] = p[hi5]
+        b[hi5] = q[hi5]
+                
+        r = r.unsqueeze(1)
+        g = g.unsqueeze(1)
+        b = b.unsqueeze(1)
+        rgb = torch.cat([r, g, b], dim=1)
+        if self.gated2:
+            rgb = rgb * self.alpha
+        return rgb

JarvisIR/package/agent_tools/HVICIDNet/net/LCA.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+from .transformer_utils import *
+
+# Cross Attention Block
+class CAB(nn.Module):
+    def __init__(self, dim, num_heads, bias):
+        super(CAB, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.q = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+        self.q_dwconv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim, bias=bias)
+        self.kv = nn.Conv2d(dim, dim*2, kernel_size=1, bias=bias)
+        self.kv_dwconv = nn.Conv2d(dim*2, dim*2, kernel_size=3, stride=1, padding=1, groups=dim*2, bias=bias)
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x, y):
+        b, c, h, w = x.shape
+
+        q = self.q_dwconv(self.q(x))
+        kv = self.kv_dwconv(self.kv(y))
+        k, v = kv.chunk(2, dim=1)
+
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = nn.functional.softmax(attn,dim=-1)
+
+        out = (attn @ v)
+
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+    
+
+# Intensity Enhancement Layer
+class IEL(nn.Module):
+    def __init__(self, dim, ffn_expansion_factor=2.66, bias=False):
+        super(IEL, self).__init__()
+
+        hidden_features = int(dim*ffn_expansion_factor)
+
+        self.project_in = nn.Conv2d(dim, hidden_features*2, kernel_size=1, bias=bias)
+        
+        self.dwconv = nn.Conv2d(hidden_features*2, hidden_features*2, kernel_size=3, stride=1, padding=1, groups=hidden_features*2, bias=bias)
+        self.dwconv1 = nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1, groups=hidden_features, bias=bias)
+        self.dwconv2 = nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1, groups=hidden_features, bias=bias)
+       
+        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+        self.Tanh = nn.Tanh()
+    def forward(self, x):
+        x = self.project_in(x)
+        x1, x2 = self.dwconv(x).chunk(2, dim=1)
+        x1 = self.Tanh(self.dwconv1(x1)) + x1
+        x2 = self.Tanh(self.dwconv2(x2)) + x2
+        x = x1 * x2
+        x = self.project_out(x)
+        return x
+  
+  
+# Lightweight Cross Attention
+class HV_LCA(nn.Module):
+    def __init__(self, dim,num_heads, bias=False):
+        super(HV_LCA, self).__init__()
+        self.gdfn = IEL(dim) # IEL and CDL have same structure
+        self.norm = LayerNorm(dim)
+        self.ffn = CAB(dim, num_heads, bias)
+        
+    def forward(self, x, y):
+        x = x + self.ffn(self.norm(x),self.norm(y))
+        x = self.gdfn(self.norm(x))
+        return x
+    
+class I_LCA(nn.Module):
+    def __init__(self, dim,num_heads, bias=False):
+        super(I_LCA, self).__init__()
+        self.norm = LayerNorm(dim)
+        self.gdfn = IEL(dim)
+        self.ffn = CAB(dim, num_heads, bias=bias)
+        
+    def forward(self, x, y):
+        x = x + self.ffn(self.norm(x),self.norm(y))
+        x = x + self.gdfn(self.norm(x)) 
+        return x

JarvisIR/package/agent_tools/HVICIDNet/net/transformer_utils.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+    with shape (batch_size, channels, height, width).
+    """
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_first"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError 
+        self.normalized_shape = (normalized_shape, )
+    
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+class NormDownsample(nn.Module):
+    def __init__(self,in_ch,out_ch,scale=0.5,use_norm=False):
+        super(NormDownsample, self).__init__()
+        self.use_norm=use_norm
+        if self.use_norm:
+            self.norm=LayerNorm(out_ch)
+        self.prelu = nn.PReLU()
+        self.down = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch,kernel_size=3,stride=1, padding=1, bias=False),
+            nn.UpsamplingBilinear2d(scale_factor=scale))
+    def forward(self, x):
+        x = self.down(x)
+        x = self.prelu(x)
+        if self.use_norm:
+            x = self.norm(x)
+            return x
+        else:
+            return x
+
+class NormUpsample(nn.Module):
+    def __init__(self, in_ch,out_ch,scale=2,use_norm=False):
+        super(NormUpsample, self).__init__()
+        self.use_norm=use_norm
+        if self.use_norm:
+            self.norm=LayerNorm(out_ch)
+        self.prelu = nn.PReLU()
+        self.up_scale = nn.Sequential(
+            nn.Conv2d(in_ch,out_ch,kernel_size=3,stride=1, padding=1, bias=False),
+            nn.UpsamplingBilinear2d(scale_factor=scale))
+        self.up = nn.Conv2d(out_ch*2,out_ch,kernel_size=1,stride=1, padding=0, bias=False)
+            
+    def forward(self, x,y):
+        x = self.up_scale(x)
+        x = torch.cat([x, y],dim=1)
+        x = self.up(x)
+        x = self.prelu(x)
+        if self.use_norm:
+            return self.norm(x)
+        else:
+            return x
+ 

JarvisIR/package/agent_tools/HVICIDNet/wavelet.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+
+
+def Normalize(x):
+    ymax = 255
+    ymin = 0
+    xmax = x.max()
+    xmin = x.min()
+    return (ymax-ymin)*(x-xmin)/(xmax-xmin) + ymin
+
+
+def dwt_init(x):
+
+    x01 = x[:, :, 0::2, :] / 2
+    x02 = x[:, :, 1::2, :] / 2
+    x1 = x01[:, :, :, 0::2]
+    x2 = x02[:, :, :, 0::2]
+    x3 = x01[:, :, :, 1::2]
+    x4 = x02[:, :, :, 1::2]
+    x_LL = x1 + x2 + x3 + x4
+    x_HL = -x1 - x2 + x3 + x4
+    x_LH = -x1 + x2 - x3 + x4
+    x_HH = x1 - x2 - x3 + x4
+
+    return torch.cat((x_LL, x_HL, x_LH, x_HH), 0)
+
+
+# 使用哈尔 haar 小波变换来实现二维离散小波
+def iwt_init(x):
+    r = 2
+    in_batch, in_channel, in_height, in_width = x.size()
+    out_batch, out_channel, out_height, out_width = int(in_batch/(r**2)),in_channel, r * in_height, r * in_width
+    x1 = x[0:out_batch, :, :] / 2
+    x2 = x[out_batch:out_batch * 2, :, :, :] / 2
+    x3 = x[out_batch * 2:out_batch * 3, :, :, :] / 2
+    x4 = x[out_batch * 3:out_batch * 4, :, :, :] / 2
+
+    h = torch.zeros([out_batch, out_channel, out_height,
+                     out_width]).float().to(x.device)
+
+    h[:, :, 0::2, 0::2] = x1 - x2 - x3 + x4
+    h[:, :, 1::2, 0::2] = x1 - x2 + x3 - x4
+    h[:, :, 0::2, 1::2] = x1 + x2 - x3 - x4
+    h[:, :, 1::2, 1::2] = x1 + x2 + x3 + x4
+
+    return h
+
+
+class DWT(nn.Module):
+    def __init__(self):
+        super(DWT, self).__init__()
+        self.requires_grad = False  # 信号处理，非卷积运算，不需要进行梯度求导
+
+    def forward(self, x):
+        return dwt_init(x)
+
+
+class IWT(nn.Module):
+    def __init__(self):
+        super(IWT, self).__init__()
+        self.requires_grad = False
+
+    def forward(self, x):
+        return iwt_init(x)

JarvisIR/package/agent_tools/IDT/analyse/cal_rf_bf.py

@@ -0,0 +1,68 @@
+import os
+import cv2
+import numpy as np
+from PIL import Image
+
+def getimgarr(path_1):
+    dir_1 = sorted(os.listdir(path_1))
+    imgs = []
+    for j in range(len(dir_1)):
+        img = cv2.imread(path_1 + '/' +dir_1[j])
+        imgs.append(img)
+    imgs = np.array(imgs, dtype=float)
+
+    return imgs
+
+#inp_path = '/disk1/release/DayRainDrop/Clear'
+#inp_path = '/disk1/release/NightRainDrop/Clear'
+#inp_path = '/disk1/release/DayRainDrop_Test/Clear'
+inp_path = '/disk1/release/NightRainDrop_Test/Clear'
+
+countall = 0
+countbf = 0
+countrf = 0
+
+sid = sorted(os.listdir(inp_path))
+print(sid,len(sid))
+
+for i in range(len(sid)):
+
+    inp_path_clean = os.path.join(inp_path, sid[i])
+    inp_path_blur = inp_path_clean.replace('/Clear/','/Blur/')
+    #print(inp_path_clean,inp_path_blur)
+
+    dropimgs = sorted(os.listdir(inp_path_clean))
+    Droplist = []
+    for didx in range(len(dropimgs)):
+        if dropimgs[didx].endswith('.png'):
+            Droplist.append(dropimgs[didx])
+    # print(Droplist)
+    # print('-------------------------------------------------')
+
+    for j in range(len(Droplist)):
+        cleaninp_frame_name = os.path.join(inp_path_clean, Droplist[j])
+        blurinp_frame_name = cleaninp_frame_name.replace('/Clear/','/Blur/')
+        # print(cleaninp_frame_name,blurinp_frame_name)
+        countall = countall + 1
+        cleanimage = Image.open(cleaninp_frame_name)
+        blurimage = Image.open(blurinp_frame_name)
+        if cleanimage == blurimage:
+            countbf = countbf + 1
+            parts = cleaninp_frame_name.split(os.sep)
+            new_path = os.path.join(parts[-2], parts[-1])
+            print(new_path)
+        else:
+            countrf = countrf + 1
+print('-countall-',countall,'-background focus-',countbf,'-raindrop focus-',countrf)
+
+#DayRainDrop         -countall- 5442,-background focus- 1836,-raindrop focus- 3606,
+#NightRainDrop       -countall- 9744,-background focus- 4906,-raindrop focus- 4838,
+#Total: 15186
+
+#DayRainDrop_Train   -countall- 8655,-background focus- 4143,-raindrop focus- 4512
+#DayRainDrop_Test    -countall- 729 -background focus- 261 -raindrop focus- 468
+#NightRainDrop_Train -countall- 4713,-background focus- 1575,-raindrop focus- 3138
+#NightRainDrop_Test  -countall- 1089 -background focus- 763 -raindrop focus- 326
+
+
+

JarvisIR/package/agent_tools/IDT/configs/daytime_128.yml

@@ -0,0 +1,55 @@
+data:
+    dataset: "RainDrop"
+    #image_size: 64 # DIT, RDiffusion, onego
+    #image_size: 256 # ICRA, Uformer, atgan
+    image_size: 128 # IDT, restormer
+    channels: 3
+    num_workers: 8
+    data_dir: "/data4/lx_workspace/datasets/SFT-data/mini/daytime_driving_rainy/rainy/"
+    conditional: True
+
+model:
+    in_channels: 3
+    out_ch: 3
+    ch: 128
+    ch_mult: [1, 1, 2, 2, 4, 4]
+    num_res_blocks: 2
+    attn_resolutions: [16, ]
+    dropout: 0.0
+    ema_rate: 0.999
+    ema: True
+    resamp_with_conv: True
+
+diffusion:
+    beta_schedule: linear
+    beta_start: 0.0001
+    beta_end: 0.02
+    num_diffusion_timesteps: 1000
+
+training:
+    #-----------DIT, RDiffusion, onego, ICRA--------------
+    # patch_n: 4
+    # batch_size: 16
+    # #---------IDT Uformer restormer--------------
+    patch_n: 1
+    batch_size: 4
+    #-----------atgan--------------
+    # patch_n: 1
+    # batch_size: 32
+    #-----------DIT, RDiffusion--------------
+    #n_epochs: 37042
+    n_epochs: 401
+    n_iters: 20000000
+    snapshot_freq: 50
+    validation_freq: 10000
+
+sampling:
+    batch_size: 4
+    last_only: True
+
+optim:
+    weight_decay: 0.000
+    optimizer: "Adam"
+    lr: 0.00002
+    amsgrad: False
+    eps: 0.00000001

JarvisIR/package/agent_tools/IDT/configs/daytime_256.yml

@@ -0,0 +1,55 @@
+data:
+    dataset: "RainDrop"
+    #image_size: 64 # DIT, RDiffusion, onego
+    image_size: 256 # ICRA, Uformer, atgan
+    # image_size: 128 # IDT, restormer
+    channels: 3
+    num_workers: 8
+    data_dir: "/data4/lx_workspace/datasets/SFT-data/mini/daytime_driving_rainy/rainy/"
+    conditional: True
+
+model:
+    in_channels: 3
+    out_ch: 3
+    ch: 128
+    ch_mult: [1, 1, 2, 2, 4, 4]
+    num_res_blocks: 2
+    attn_resolutions: [16, ]
+    dropout: 0.0
+    ema_rate: 0.999
+    ema: True
+    resamp_with_conv: True
+
+diffusion:
+    beta_schedule: linear
+    beta_start: 0.0001
+    beta_end: 0.02
+    num_diffusion_timesteps: 1000
+
+training:
+    #-----------DIT, RDiffusion, onego, ICRA--------------
+    # patch_n: 4
+    # batch_size: 16
+    # #---------IDT Uformer restormer--------------
+    patch_n: 1
+    batch_size: 4
+    #-----------atgan--------------
+    # patch_n: 1
+    # batch_size: 32
+    #-----------DIT, RDiffusion--------------
+    #n_epochs: 37042
+    n_epochs: 401
+    n_iters: 20000000
+    snapshot_freq: 50
+    validation_freq: 10000
+
+sampling:
+    batch_size: 4
+    last_only: True
+
+optim:
+    weight_decay: 0.000
+    optimizer: "Adam"
+    lr: 0.00002
+    amsgrad: False
+    eps: 0.00000001