Colab/Real-CUGAN.md

Real-CUGAN images upscaler

• Step 1: Preparation

Before you begin, make sure that you have set the runtime type to GPU (Hardware acclerator: GPU).

ROOTPATH="/content/ailab/" # root dir (a constant)
ModelPath=ROOTPATH+"Real-CUGAN/model/" # model dir
PendingPath=ROOTPATH+"Real-CUGAN/pending/" # input dir
FinishPath=ROOTPATH+"Real-CUGAN/finish/" # output dir
ModelName="up2x-latest-no-denoise.pth" # default model
Tile=4 #{0,1,2,3,4,auto}; the larger the number, the smaller the memory consumption
# initialize environment
!pip install torch opencv-python
!git clone https://github.com/bilibili/ailab.git
from google.colab import drive
drive.mount('/content/gdrive')

• Step 2: Download

Download model files from here and save them - in a folder called updated_weights - under your Google Drive's root folder.

• Step 3: upload

Put/upload your image(s) under /content/aliab/Real-CUGAN/pending.

• Step 4: Setup

Run the following and choose the model you want.

import os
if not os.path.exists(ModelPath):
    os.mkdir(ModelPath)
if not os.path.exists(PendingPath):
    os.mkdir(PendingPath)
if not os.path.exists(FinishPath):
    os.mkdir(FinishPath)
!cp -r /content/gdrive/MyDrive/updated_weights/* /content/ailab/Real-CUGAN/model/
fileNames = os.listdir(PendingPath)
print("Pending images:")
for i in fileNames:
  print("\t"+i)
fileNames = os.listdir(ModelPath)
print("Model files available:")
for idx, i in enumerate(fileNames):
  print(f"{idx+1}. \t {i}")
choice = int(input("Select model (leave blank for default): "))
if choice:
  ModelName=fileNames[choice-1]
Amplification=ModelName[2] # amplifying ratio
if (not os.path.isfile(ModelPath+ModelName)):
  print("Warning: selected model file does not exist")

• Execution

Run the processing script.

import shutil
import sys
sys.path.append("/content/ailab/Real-CUGAN")
import os
source_path = '/content/gdrive/MyDrive/upscale'
destination_path = '/content/ailab/Real-CUGAN'
os.system(f'cp -v {source_path}/pending/*.png {destination_path}/pending')
import matplotlib.pyplot as plt
%matplotlib inline
import torch
from torch import nn as nn
from torch.nn import functional as F
import cv2
import numpy as np
from upcunet_v3 import RealWaifuUpScaler
from time import time as ttime  
fileNames = os.listdir(destination_path + '/pending')
print(f"using model {ModelPath+ModelName}")
upscaler = RealWaifuUpScaler(2, ModelPath+ModelName, half=True, device="cuda:0")
t0 = ttime()
for i in fileNames:
    torch.cuda.empty_cache()
    try:
        img = cv2.imread(destination_path + '/pending/' + i)[:, :, [2, 1, 0]]
        result = upscaler(img, tile_mode=5, cache_mode=2, alpha=1)
        # Estrai il nome del file senza estensione
        file_name, file_extension = os.path.splitext(i)
        # Rinomina l'immagine con il nome del modello scelto
        new_name = f"{file_name}_{ModelName}{file_extension}"
        cv2.imwrite(os.path.join(destination_path, 'finish', 'Real_CUGAN', new_name), result[:, :, ::-1])
        # Sposta l'immagine rinominata nella cartella di destinazione
        source_file = os.path.join(destination_path, 'finish', 'Real_CUGAN', new_name)
        destination_file = os.path.join(destination_path, 'finish', 'Real_CUGAN', new_name)
        try:
            os.rename(source_file, destination_file)
            print(f'Moved {source_file} to {destination_file}')
        except OSError as e:
            print(f'Error moving {source_file} to {destination_file}: {e}')
        # Sposta anche l'immagine dalla cartella "pending" alla cartella "done"
        source_pending_file = os.path.join(source_path, 'pending', i)
        destination_done_file = os.path.join(source_path, 'done', i)
        try:
            os.rename(source_pending_file, destination_done_file)
            print(f'Moved {source_pending_file} to {destination_done_file}')
        except OSError as e:
            print(f'Error moving {source_pending_file} to {destination_done_file}: {e}')
    except RuntimeError as e:
        print (i+" FAILED")
        print (e)
    else:
        print(i+" DONE")
t1 = ttime()
print("time_spent", t1 - t0)