# Super Resolution

# Copyright (C) 2024-2026 Collabora Ltd.

#

# This library is free software; you can redistribute it and/or

# modify it under the terms of the GNU Library General Public

# License as published by the Free Software Foundation; either

# version 2 of the License, or (at your option) any later version.

#

# This library is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU

# Library General Public License for more details.

#

# You should have received a copy of the GNU Library General Public

# License along with this library; if not, write to the

# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,

# Boston, MA 02110-1301, USA.

from log.global_logger import GlobalLogger

import backend

CAN_REGISTER_ELEMENT = True

try:

from video_transform import VideoTransform

from utils.format_converter import FormatConverter

from engine.super_res_engine import ExportedSuperRes, SuperResEngine

from engine.engine_factory import EngineFactory

from backend import frameio

from tasks.superres import SuperResTask

except ImportError as e:

CAN_REGISTER_ELEMENT = False

GlobalLogger().warning(f"The 'superres' element will not be available. Error {e}")

DEFAULT_MODEL_NAME = "real-esrgan-x4"

class SuperResTransform(VideoTransform, SuperResTask):

"""

GStreamer element for image super-resolution using Real-ESRGAN.

Set model-name to a Real-ESRGAN variant: real-esrgan-x4 or real-esrgan-x2.

The upscaled frame is resized back to the original buffer dimensions

(in-place transform). This provides enhanced detail while maintaining

pipeline compatibility. Use frame-stride to reduce compute load.

"""

__gstmetadata__ = (

"Super Resolution",

"Transform",

"Image super-resolution using Real-ESRGAN",

"Aaron Boxer <aaron.boxer@collabora.com>",

)

ACCEPTS_DEVICE_FRAMES = True

def __init__(self):

super().__init__()

self.task_engine_name = self.mgr.engine_name = "pyml_superres_engine"

EngineFactory.register(self.mgr.engine_name, SuperResEngine)

self.format_converter = FormatConverter()

self.model_name = DEFAULT_MODEL_NAME

def export_model(self, model_name):

return ExportedSuperRes(model_name, self.mgr.engine_name)

def process_frames(self, frames, num_sources, fmt, target):

"""Upscale the primary frame and write it back at the original size."""

frame = frames[0] if frames.ndim == 4 else frames

upscaled = self.forward(frame)

if upscaled is None:

return

output, _blob = self.decode(upscaled, fmt)

frameio.write_result(target, output)

if CAN_REGISTER_ELEMENT and backend.BACKEND == "gst":

__gstelementfactory__ = backend.register_gst_element(

"pyml_superres", SuperResTransform

)

elif not CAN_REGISTER_ELEMENT:

GlobalLogger().warning(

"The 'pyml_superres' element will not be registered because required modules are missing."

)