filters and yolo

MANIFEST.in

@@ -2,4 +2,5 @@ include gimpml/plugins/colorpalette/color_palette.png
 include gimpml/plugins/images/plugin_logo.png
 include gimpml/plugins/images/error_icon.png
 include gimpml/tools/model_info.csv
-
+include gimpml/tools/pytorch-YOLOv4/yolo_data/coco.names
+include gimpml/tools/pytorch-YOLOv4/yolo_data/voc.names

gimpml/__init__.py

@@ -5,7 +5,7 @@ from .tools.dehaze import get_dehaze as dehaze
 from .tools.denoise import get_denoise as denoise
 from .tools.matting import get_matting as matting
 from .tools.enlighten import get_enlighten as enlighten
-
+from .tools.canny import get_edge as edge
 # from .facegen import get_newface as newface
 from .tools.faceparse import get_face as parseface
 from .tools.interpolation import get_inter as interpolateframe
@@ -14,5 +14,6 @@ from .tools.complete_install import setup_python_weights
 from .tools.semseg import get_seg as semseg
 from .tools.superresolution import get_super as super
 from .tools.inpainting import get_inpaint as inpaint
-
+from .tools.detectobjects import get_detect_objects as detect_objects
+from .filters import *
 # from .plugins.plugin_utils import show_error_dialog

gimpml/filters/__init__.py

@@ -0,0 +1,64 @@
+import os
+import pickle
+from gimpml.tools.tools_utils import get_weight_path
+import subprocess
+import numpy as np
+import cv2
+
+
+# procedure="fog"
+# image_path = r"D:\PycharmProjects\test_aug\dp.png"
+# save_path = r"D:\PycharmProjects\test_aug\dp2.png"
+# opacity = 100
+# turbulence = 1  # between o and 1
+# rgb = (244, 0, 0)
+# image = cv2.imread(image_path)[:, :, ::-1]
+# kwargs = {"image": image, "opacity": opacity, "rgb": rgb, "turbulence": turbulence}
+
+def run(procedure, **kwargs):
+    weight_path = get_weight_path()
+    if "image_path" in kwargs.keys() and "image" in kwargs.keys():
+        raise Exception("Both image_path and image should not be passed as input.")
+    if "image_path" in kwargs.keys() and not os.path.isfile(kwargs['image_path']):
+        raise Exception("Input image file does not exist.")
+    if "image" in kwargs.keys() and not isinstance(kwargs['image'], np.ndarray) and not len(kwargs['image'].shape) == 3:
+        raise Exception("Invalid input image.")
+    return_image, remove_input_image = False, False
+    if "save_path" not in kwargs.keys():
+        kwargs["save_path"] = os.path.join(weight_path, "..", "tmp_filter2.png")
+        return_image = True
+    if "image" in kwargs.keys():
+        image_path = os.path.join(weight_path, "..", "tmp_filter.png")
+        channels = kwargs["image"].shape[2]
+        if channels == 3:
+            cv2.imwrite(image_path, kwargs["image"][:, :, ::-1])
+            kwargs["image_path"] = image_path
+            kwargs.pop("image")
+            remove_input_image = True
+            # print("Image saved.")
+        elif channels == 4:
+            cv2.imwrite(image_path, kwargs["image"][:, :, [2, 1, 0, 3]])
+            kwargs["image_path"] = image_path
+            kwargs.pop("image")
+            remove_input_image = True
+            # print("Image saved.")
+        else:
+            raise Exception("High-dimensional image not supported.")
+    with open(os.path.join(weight_path, "..", "gimp_ml_augment.pkl"), "wb") as file:
+        pickle.dump(
+            kwargs,
+            file,
+        )
+    command_str = "gimp-2.99 -idf --batch-interpreter=python-fu-eval -b - < " + os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), procedure.strip()
+    ) + ".py"
+    # print(command_str)
+    subprocess.call(command_str, shell=True)
+    if remove_input_image:
+        os.remove(kwargs["image_path"])
+    if return_image:
+        img = cv2.imread(kwargs["save_path"])
+        channels = img.shape[2]
+        os.remove(kwargs["save_path"])
+        return img[:, :, ::-1] if channels == 3 else img[:, :, [2, 1, 0, 3]]
+

gimpml/filters/fog.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+# coding: utf-8
+try:
+    import time
+    import os
+    import pickle
+    import gi
+
+    gi.require_version("Gimp", "3.0")
+    from gi.repository import Gimp, GObject, Gio
+
+    install_location = os.path.join(os.path.expanduser("~"), "GIMP-ML")
+    with open(os.path.join(install_location, "gimp_ml_augment.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+
+    image_path = data_output['image_path']
+    opacity = data_output['opacity']
+    rgb = data_output['rgb']
+    save_path = data_output['save_path']
+    turbulence = data_output['turbulence']
+
+    image = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(image_path))  # image
+    image_layer = image.get_active_layer()
+
+    if image.get_base_type() is Gimp.ImageBaseType.RGB:
+        type = Gimp.ImageType.RGBA_IMAGE
+    else:
+        type = Gimp.ImageType.GRAYA_IMAGE
+
+    # fog = Gimp.Layer.new_from_drawable(image_layer, image)
+    fog = Gimp.Layer.new(image, "tmp",
+                         image_layer.get_width(), image_layer.get_height(),
+                         type, opacity,
+                         Gimp.LayerMode.NORMAL)
+    fog.fill(Gimp.FillType.TRANSPARENT)
+    image.insert_layer(fog, image_layer.get_parent(), image.get_item_position(image_layer))
+    color = Gimp.RGB()
+    color.set(rgb[0], rgb[1], rgb[2])
+    Gimp.context_set_background(color)
+    fog.edit_fill(Gimp.FillType.BACKGROUND)
+
+    # create a layer mask for the new layer
+    mask = fog.create_mask(0)
+    fog.add_mask(mask)
+
+    # add some clouds to the layer
+    Gimp.get_pdb().run_procedure('plug-in-plasma', [
+        GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
+        GObject.Value(Gimp.Image, image),
+        GObject.Value(Gimp.Drawable, mask),
+        GObject.Value(GObject.TYPE_INT, int(time.time())),
+        GObject.Value(GObject.TYPE_DOUBLE, turbulence),
+    ])
+    # apply the clouds to the layer
+    fog.remove_mask(Gimp.MaskApplyMode.APPLY)
+    fog.set_visible(True)
+
+    thumb = image.duplicate()
+    layer = thumb.merge_visible_layers(Gimp.MergeType.CLIP_TO_IMAGE)
+
+    # save
+    interlace, compression = 0, 2
+    Gimp.get_pdb().run_procedure(
+        "file-png-save",
+        [
+            GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
+            GObject.Value(Gimp.Image, thumb),
+            GObject.Value(GObject.TYPE_INT, 1),
+            GObject.Value(
+                Gimp.ObjectArray, Gimp.ObjectArray.new(Gimp.Drawable, [layer], 0)
+            ),
+            GObject.Value(
+                Gio.File,
+                Gio.File.new_for_path(save_path),
+            ),
+            GObject.Value(GObject.TYPE_BOOLEAN, interlace),
+            GObject.Value(GObject.TYPE_INT, compression),
+            # write all PNG chunks except oFFs(ets)
+            GObject.Value(GObject.TYPE_BOOLEAN, True),
+            GObject.Value(GObject.TYPE_BOOLEAN, True),
+            GObject.Value(GObject.TYPE_BOOLEAN, False),
+            GObject.Value(GObject.TYPE_BOOLEAN, True),
+        ],
+    )
+
+    # quit
+    Gimp.get_pdb().run_procedure("gimp-quit", [GObject.Value(GObject.TYPE_BOOLEAN, True)])
+
+except:
+    pass

gimpml/filters/gaussian_blur.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+# coding: utf-8
+
+# def add_gaussian_blur(image_path, save_path):
+try:
+    import os
+    import pickle
+    import gi
+    gi.require_version("Gimp", "3.0")
+    from gi.repository import Gimp, GObject, Gio
+
+    install_location = os.path.join(os.path.expanduser("~"), "GIMP-ML")
+    with open(os.path.join(install_location, "gimp_ml_augment.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+
+    image_path = data_output['image_path']
+    save_path = data_output['save_path']
+    horizontal = data_output['horizontal']
+    vertical = data_output['vertical']
+    method = data_output['method']
+
+    image = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(image_path))  # image
+    image_layer = image.get_active_layer()  # drawable
+
+    # run plugin
+    Gimp.get_pdb().run_procedure(
+        "plug-in-gauss",
+        [
+            GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
+            GObject.Value(Gimp.Image, image),
+            GObject.Value(Gimp.Drawable, image_layer),
+            GObject.Value(GObject.TYPE_DOUBLE, horizontal),
+            GObject.Value(GObject.TYPE_DOUBLE, vertical),
+            GObject.Value(GObject.TYPE_INT, method),
+        ]
+    )
+    image_layer = image.get_active_layer()
+
+    # save
+    interlace, compression = 0, 2
+    Gimp.get_pdb().run_procedure(
+        "file-png-save",
+        [
+            GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
+            GObject.Value(Gimp.Image, image),
+            GObject.Value(GObject.TYPE_INT, 1),
+            GObject.Value(
+                Gimp.ObjectArray, Gimp.ObjectArray.new(Gimp.Drawable, [image_layer], 0)
+            ),
+            GObject.Value(
+                Gio.File,
+                Gio.File.new_for_path(save_path),
+            ),
+            GObject.Value(GObject.TYPE_BOOLEAN, interlace),
+            GObject.Value(GObject.TYPE_INT, compression),
+            # write all PNG chunks except oFFs(ets)
+            GObject.Value(GObject.TYPE_BOOLEAN, True),
+            GObject.Value(GObject.TYPE_BOOLEAN, True),
+            GObject.Value(GObject.TYPE_BOOLEAN, False),
+            GObject.Value(GObject.TYPE_BOOLEAN, True),
+        ],
+    )
+
+    # quit
+    Gimp.get_pdb().run_procedure("gimp-quit", [GObject.Value(GObject.TYPE_BOOLEAN, True)])
+
+except:
+    pass
+# gimp-2.99 -idf --batch-interpreter=python-fu-eval -b - < gaussian_blur.py

gimpml/plugins/canny/canny.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+# coding: utf-8
+"""
+ .d8888b.  8888888 888b     d888 8888888b.       888b     d888 888
+d88P  Y88b   888   8888b   d8888 888   Y88b      8888b   d8888 888
+888    888   888   88888b.d88888 888    888      88888b.d88888 888
+888          888   888Y88888P888 888   d88P      888Y88888P888 888
+888  88888   888   888 Y888P 888 8888888P"       888 Y888P 888 888
+888    888   888   888  Y8P  888 888             888  Y8P  888 888
+Y88b  d88P   888   888   "   888 888             888   "   888 888
+ "Y8888P88 8888888 888       888 888             888       888 88888888
+
+Performs Canny Edge Detection for current layer.
+"""
+import gi
+gi.require_version("Gimp", "3.0")
+gi.require_version("GimpUi", "3.0")
+gi.require_version("Gtk", "3.0")
+from gi.repository import Gimp, GimpUi, GObject, GLib, Gio, Gtk
+import gettext
+import subprocess
+import pickle
+import os
+import sys
+sys.path.extend([os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")])
+from plugin_utils import *
+
+
+_ = gettext.gettext
+image_paths = {
+    "colorpalette": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "..",
+        "colorpalette",
+        "color_palette.png",
+    ),
+    "logo": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "images", "plugin_logo.png"
+    ),
+    "error": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "images", "error_icon.png"
+    ),
+}
+
+
+def canny(
+    procedure, image, drawable, min_val, max_val, progress_bar, config_path_output
+):
+    # Save inference parameters and layers
+    weight_path = config_path_output["weight_path"]
+    python_path = config_path_output["python_path"]
+    plugin_path = config_path_output["plugin_path"]
+
+    Gimp.context_push()
+    image.undo_group_start()
+
+    save_image(image, drawable, os.path.join(weight_path, "..", "cache.png"))
+
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+        pickle.dump(
+            {
+                "min_val": int(min_val),
+                "max_val": int(max_val),
+                "inference_status": "started",
+            },
+            file,
+        )
+
+    # Run inference and load as layer
+    subprocess.call([python_path, plugin_path])
+    subprocess.call([python_path, plugin_path])
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+    image.undo_group_end()
+    Gimp.context_pop()
+    if data_output["inference_status"] == "success":
+        result = Gimp.file_load(
+            Gimp.RunMode.NONINTERACTIVE,
+            Gio.file_new_for_path(os.path.join(weight_path, "..", "cache.png")),
+        )
+        result_layer = result.get_active_layer()
+        copy = Gimp.Layer.new_from_drawable(result_layer, image)
+        copy.set_name("Canny Edge")
+        copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)  # DIFFERENCE_LEGACY
+        image.insert_layer(copy, None, -1)
+
+        # Remove temporary layers that were saved
+        my_dir = os.path.join(weight_path, "..")
+        for f_name in os.listdir(my_dir):
+            if f_name.startswith("cache"):
+                os.remove(os.path.join(my_dir, f_name))
+
+        return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
+    else:
+        show_dialog(
+            "Inference not successful. See error_log.txt in GIMP-ML folder.",
+            "Error !",
+            "error",
+            image_paths
+        )
+        return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
+
+
+def run(procedure, run_mode, image, n_drawables, layer, args, data):
+    min_val = args.index(0)
+    max_val = args.index(1)
+
+    if run_mode == Gimp.RunMode.INTERACTIVE:
+        # Get all paths
+        config_path = os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools"
+        )
+        with open(os.path.join(config_path, "gimp_ml_config.pkl"), "rb") as file:
+            config_path_output = pickle.load(file)
+        python_path = config_path_output["python_path"]
+        config_path_output["plugin_path"] = os.path.join(config_path, "canny.py")
+
+        config = procedure.create_config()
+        config.begin_run(image, run_mode, args)
+
+        GimpUi.init("canny.py")
+        use_header_bar = Gtk.Settings.get_default().get_property(
+            "gtk-dialogs-use-header"
+        )
+        dialog = GimpUi.Dialog(use_header_bar=use_header_bar, title=_("Canny Edge..."))
+        dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
+        dialog.add_button("_Help", Gtk.ResponseType.APPLY)
+        dialog.add_button("_Run Inference", Gtk.ResponseType.OK)
+
+        vbox = Gtk.Box(
+            orientation=Gtk.Orientation.VERTICAL, homogeneous=False, spacing=10
+        )
+        dialog.get_content_area().add(vbox)
+        vbox.show()
+
+        # Create grid to set all the properties inside.
+        grid = Gtk.Grid()
+        grid.set_column_homogeneous(False)
+        grid.set_border_width(10)
+        grid.set_column_spacing(10)
+        grid.set_row_spacing(10)
+        vbox.add(grid)
+        grid.show()
+
+        # Show Logo
+        logo = Gtk.Image.new_from_file(image_paths["logo"])
+        # grid.attach(logo, 0, 0, 1, 1)
+        vbox.pack_start(logo, False, False, 1)
+        logo.show()
+
+        # Show License
+        license_text = _("PLUGIN LICENSE : Apache 2")
+        label = Gtk.Label(label=license_text)
+        # grid.attach(label, 1, 1, 1, 1)
+        vbox.pack_start(label, False, False, 1)
+        label.show()
+
+        # min_val parameter
+        label = Gtk.Label.new_with_mnemonic(_("_Min"))
+        grid.attach(label, 0, 0, 1, 1)
+        label.show()
+        spin = GimpUi.prop_spin_button_new(
+            config, "min_val", step_increment=1, page_increment=10, digits=0
+        )
+        grid.attach(spin, 1, 0, 1, 1)
+        spin.show()
+
+        # max_val parameter
+        label = Gtk.Label.new_with_mnemonic(_("_Max"))
+        grid.attach(label, 2, 0, 1, 1)
+        label.show()
+        spin = GimpUi.prop_spin_button_new(
+            config, "max_val", step_increment=1, page_increment=10, digits=0
+        )
+        grid.attach(spin, 3, 0, 1, 1)
+        spin.show()
+
+        progress_bar = Gtk.ProgressBar()
+        vbox.add(progress_bar)
+        progress_bar.show()
+
+        # Wait for user to click
+        dialog.show()
+        while True:
+            response = dialog.run()
+            if response == Gtk.ResponseType.OK:
+                min_val = config.get_property("min_val")
+                max_val = config.get_property("max_val")
+                result = canny(
+                    procedure,
+                    image,
+                    layer,
+                    min_val,
+                    max_val,
+                    progress_bar,
+                    config_path_output,
+                )
+                # If the execution was successful, save parameters so they will be restored next time we show dialog.
+                if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
+                    config.end_run(Gimp.PDBStatusType.SUCCESS)
+                return result
+            elif response == Gtk.ResponseType.APPLY:
+                url = "https://kritiksoman.github.io/GIMP-ML-Docs/docs-page.html#item-7-16"
+                Gio.app_info_launch_default_for_uri(url, None)
+                continue
+            else:
+                dialog.destroy()
+                return procedure.new_return_values(
+                    Gimp.PDBStatusType.CANCEL, GLib.Error()
+                )
+
+
+class Canny(Gimp.PlugIn):
+    ## Parameters ##
+    __gproperties__ = {
+        "min_val": (
+            float,
+            _("_Min"),
+            "Min threshold",
+            0,
+            254,
+            100,
+            GObject.ParamFlags.READWRITE,
+        ),
+        "max_val": (
+            float,
+            _("_Max"),
+            "Max threshold",
+            1,
+            255,
+            200,
+            GObject.ParamFlags.READWRITE,
+        ),
+    }
+
+    ## GimpPlugIn virtual methods ##
+    def do_query_procedures(self):
+        self.set_translation_domain(
+            "gimp30-python", Gio.file_new_for_path(Gimp.locale_directory())
+        )
+        return ["canny"]
+
+    def do_create_procedure(self, name):
+        procedure = None
+        if name == "canny":
+            procedure = Gimp.ImageProcedure.new(
+                self, name, Gimp.PDBProcType.PLUGIN, run, None
+            )
+
+            procedure.set_image_types("*")
+            procedure.set_documentation(
+                N_("Performs Canny Edge Detection for current layer."),
+                globals()[
+                    "__doc__"
+                ],  # This includes the docstring, on the top of the file
+                name,
+            )
+            procedure.set_menu_label(N_("_Canny Edge..."))
+            procedure.set_attribution("Kritik Soman", "GIMP-ML", "2021")
+            procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
+            procedure.add_argument_from_property(self, "min_val")
+            procedure.add_argument_from_property(self, "max_val")
+
+        return procedure
+
+
+Gimp.main(Canny.__gtype__, sys.argv)

gimpml/plugins/detectobjects/detectobjects.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+# coding: utf-8
+"""
+ .d8888b.  8888888 888b     d888 8888888b.       888b     d888 888
+d88P  Y88b   888   8888b   d8888 888   Y88b      8888b   d8888 888
+888    888   888   88888b.d88888 888    888      88888b.d88888 888
+888          888   888Y88888P888 888   d88P      888Y88888P888 888
+888  88888   888   888 Y888P 888 8888888P"       888 Y888P 888 888
+888    888   888   888  Y8P  888 888             888  Y8P  888 888
+Y88b  d88P   888   888   "   888 888             888   "   888 888
+ "Y8888P88 8888888 888       888 888             888       888 88888888
+
+
+Object detection on the current layer.
+"""
+import gi
+gi.require_version("Gimp", "3.0")
+gi.require_version("GimpUi", "3.0")
+gi.require_version("Gtk", "3.0")
+from gi.repository import Gimp, GimpUi, GObject, GLib, Gio, Gtk
+import gettext
+import subprocess
+import pickle
+import os
+import sys
+sys.path.extend([os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")])
+from plugin_utils import *
+
+
+_ = gettext.gettext
+image_paths = {
+    "logo": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "images", "plugin_logo.png"
+    ),
+    "error": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "images", "error_icon.png"
+    ),
+}
+
+
+def detectobjects(procedure, image, drawable, force_cpu, progress_bar, config_path_output):
+    # Save inference parameters and layers
+    weight_path = config_path_output["weight_path"]
+    python_path = config_path_output["python_path"]
+    plugin_path = config_path_output["plugin_path"]
+
+    Gimp.context_push()
+    image.undo_group_start()
+
+    save_image(image, drawable, os.path.join(weight_path, "..", "cache.png"))
+
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+        pickle.dump({"force_cpu": bool(force_cpu), "inference_status": "started", "get_predict_image": True}, file)
+    image.undo_group_end()
+    Gimp.context_pop()
+    # Run inference and load as layer
+    subprocess.call([python_path, plugin_path])
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+    if data_output["inference_status"] == "success":
+        result = Gimp.file_load(
+            Gimp.RunMode.NONINTERACTIVE,
+            Gio.file_new_for_path(os.path.join(weight_path, "..", "cache.png")),
+        )
+        result_layer = result.get_active_layer()
+        copy = Gimp.Layer.new_from_drawable(result_layer, image)
+        copy.set_name("Objects detected.")
+        copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)  # DIFFERENCE_LEGACY
+        image.insert_layer(copy, None, -1)
+
+        # Remove temporary layers that were saved
+        my_dir = os.path.join(weight_path, "..")
+        for f_name in os.listdir(my_dir):
+            if f_name.startswith("cache"):
+                os.remove(os.path.join(my_dir, f_name))
+
+        return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
+
+    else:
+        show_dialog(
+            "Inference not successful. See error_log.txt in GIMP-ML folder.",
+            "Error !",
+            "error",
+            image_paths
+        )
+        return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
+
+
+def run(procedure, run_mode, image, n_drawables, layer, args, data):
+    force_cpu = args.index(0)
+
+    if run_mode == Gimp.RunMode.INTERACTIVE:
+        # Get all paths
+        config_path = os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools"
+        )
+        with open(os.path.join(config_path, "gimp_ml_config.pkl"), "rb") as file:
+            config_path_output = pickle.load(file)
+        python_path = config_path_output["python_path"]
+        config_path_output["plugin_path"] = os.path.join(config_path, "detectobjects.py")
+
+        config = procedure.create_config()
+        config.set_property("force_cpu", force_cpu)
+        config.begin_run(image, run_mode, args)
+
+        GimpUi.init("detectobjects.py")
+        use_header_bar = Gtk.Settings.get_default().get_property(
+            "gtk-dialogs-use-header"
+        )
+        dialog = GimpUi.Dialog(use_header_bar=use_header_bar, title=_("Detect Objects..."))
+        dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
+        dialog.add_button("_Help", Gtk.ResponseType.APPLY)
+        dialog.add_button("_Run Inference", Gtk.ResponseType.OK)
+
+        vbox = Gtk.Box(
+            orientation=Gtk.Orientation.VERTICAL, homogeneous=False, spacing=10
+        )
+        dialog.get_content_area().add(vbox)
+        vbox.show()
+
+        # Create grid to set all the properties inside.
+        grid = Gtk.Grid()
+        grid.set_column_homogeneous(False)
+        grid.set_border_width(10)
+        grid.set_column_spacing(10)
+        grid.set_row_spacing(10)
+        vbox.add(grid)
+        grid.show()
+
+        # Force CPU parameter
+        spin = GimpUi.prop_check_button_new(config, "force_cpu", _("Force _CPU"))
+        spin.set_tooltip_text(
+            _(
+                "If checked, CPU is used for model inference."
+                " Otherwise, GPU will be used if available."
+            )
+        )
+        grid.attach(spin, 1, 2, 1, 1)
+        spin.show()
+
+        # Show Logo
+        logo = Gtk.Image.new_from_file(image_paths["logo"])
+        # grid.attach(logo, 0, 0, 1, 1)
+        vbox.pack_start(logo, False, False, 1)
+        logo.show()
+
+        # Show License
+        license_text = _("PLUGIN LICENSE : Apache-2.0")
+        label = Gtk.Label(label=license_text)
+        # grid.attach(label, 1, 1, 1, 1)
+        vbox.pack_start(label, False, False, 1)
+        label.show()
+
+        progress_bar = Gtk.ProgressBar()
+        vbox.add(progress_bar)
+        progress_bar.show()
+
+        # Wait for user to click
+        dialog.show()
+        while True:
+            response = dialog.run()
+            if response == Gtk.ResponseType.OK:
+                force_cpu = config.get_property("force_cpu")
+                result = detectobjects(
+                    procedure, image, layer, force_cpu, progress_bar, config_path_output
+                )
+                # super_resolution(procedure, image, n_drawables, layer, force_cpu, progress_bar, config_path_output)
+                # If the execution was successful, save parameters so they will be restored next time we show dialog.
+                if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
+                    config.end_run(Gimp.PDBStatusType.SUCCESS)
+                return result
+            elif response == Gtk.ResponseType.APPLY:
+                url = "https://kritiksoman.github.io/GIMP-ML-Docs/docs-page.html#item-7-14"
+                Gio.app_info_launch_default_for_uri(url, None)
+                continue
+            else:
+                dialog.destroy()
+                return procedure.new_return_values(
+                    Gimp.PDBStatusType.CANCEL, GLib.Error()
+                )
+
+
+class DetectObjects(Gimp.PlugIn):
+    ## Parameters ##
+    __gproperties__ = {
+        "force_cpu": (
+            bool,
+            _("Force _CPU"),
+            "Force CPU",
+            False,
+            GObject.ParamFlags.READWRITE,
+        ),
+    }
+
+    ## GimpPlugIn virtual methods ##
+    def do_query_procedures(self):
+        self.set_translation_domain(
+            "gimp30-python", Gio.file_new_for_path(Gimp.locale_directory())
+        )
+        return ["detectobjects"]
+
+    def do_create_procedure(self, name):
+        procedure = None
+        if name == "detectobjects":
+            procedure = Gimp.ImageProcedure.new(
+                self, name, Gimp.PDBProcType.PLUGIN, run, None
+            )
+            procedure.set_image_types("*")
+            procedure.set_documentation(
+                N_("Detects objects on the current layer."),
+                globals()[
+                    "__doc__"
+                ],  # This includes the docstring, on the top of the file
+                name,
+            )
+            procedure.set_menu_label(N_("Detect Objects..."))
+            procedure.set_attribution("Kritik Soman", "GIMP-ML", "2021")
+            procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
+            procedure.add_argument_from_property(self, "force_cpu")
+
+        return procedure
+
+
+Gimp.main(DetectObjects.__gtype__, sys.argv)

gimpml/plugins/filterfolder/filterfolder.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python3
+# coding: utf-8
+"""
+ .d8888b.  8888888 888b     d888 8888888b.       888b     d888 888
+d88P  Y88b   888   8888b   d8888 888   Y88b      8888b   d8888 888
+888    888   888   88888b.d88888 888    888      88888b.d88888 888
+888          888   888Y88888P888 888   d88P      888Y88888P888 888
+888  88888   888   888 Y888P 888 8888888P"       888 Y888P 888 888
+888    888   888   888  Y8P  888 888             888  Y8P  888 888
+Y88b  d88P   888   888   "   888 888             888   "   888 888
+ "Y8888P88 8888888 888       888 888             888       888 88888888
+
+
+Object detection on the current layer.
+"""
+import gi
+
+gi.require_version("Gimp", "3.0")
+gi.require_version("GimpUi", "3.0")
+gi.require_version("Gtk", "3.0")
+from gi.repository import Gimp, GimpUi, GObject, GLib, Gio, Gtk
+import gettext
+import subprocess
+import pickle
+import os
+import sys
+
+sys.path.extend([os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")])
+from plugin_utils import *
+
+_ = gettext.gettext
+image_paths = {
+    "logo": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "images", "plugin_logo.png"
+    ),
+    "error": os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "images", "error_icon.png"
+    ),
+}
+
+
+def filterfolder(procedure, args_dict, config_path_output):
+    # Save inference parameters and layers
+    weight_path = config_path_output["weight_path"]
+    python_path = config_path_output["python_path"]
+    plugin_path = config_path_output["plugin_path"]
+    #
+    # Gimp.context_push()
+    # image.undo_group_start()
+    #
+    # save_image(image, drawable, os.path.join(weight_path, "..", "cache.png"))
+    args_dict["inference_status"] = "started"
+    args_dict["get_predict_image"] = False
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+        pickle.dump(args_dict, file)
+    # image.undo_group_end()
+    # Gimp.context_pop()
+    # Run inference and load as layer
+    subprocess.call([python_path, plugin_path])
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+    if data_output["inference_status"] == "success":
+        count = data_output["count"]
+        if count == 0:
+            message = "No files found with entered objects."
+        elif count == 1:
+            message = str(count) + " file moved to " + os.path.join(args_dict["image_path"], "filtered")
+        else:
+            message = str(count) + " files moved to " + os.path.join(args_dict["image_path"], "filtered")
+        show_dialog(
+            message,
+            "Inference Complete.",
+            "logo",
+            image_paths
+        )
+        # result = Gimp.file_load(
+        #     Gimp.RunMode.NONINTERACTIVE,
+        #     Gio.file_new_for_path(os.path.join(weight_path, "..", "cache.png")),
+        # )
+        # result_layer = result.get_active_layer()
+        # copy = Gimp.Layer.new_from_drawable(result_layer, image)
+        # copy.set_name("Objects detected.")
+        # copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)  # DIFFERENCE_LEGACY
+        # image.insert_layer(copy, None, -1)
+        #
+        # # Remove temporary layers that were saved
+        # my_dir = os.path.join(weight_path, "..")
+        # for f_name in os.listdir(my_dir):
+        #     if f_name.startswith("cache"):
+        #         os.remove(os.path.join(my_dir, f_name))
+
+        return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
+
+    else:
+        show_dialog(
+            "Inference not successful. See error_log.txt in GIMP-ML folder.",
+            "Error !",
+            "error",
+            image_paths
+        )
+        return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
+
+
+def run(procedure, run_mode, args):
+    args_dict = {}
+    # force_cpu = args.index(0)
+
+    # if run_mode == Gimp.RunMode.INTERACTIVE:
+    # Get all paths
+    config_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools"
+    )
+    with open(os.path.join(config_path, "gimp_ml_config.pkl"), "rb") as file:
+        config_path_output = pickle.load(file)
+    python_path = config_path_output["python_path"]
+    config_path_output["plugin_path"] = os.path.join(config_path, "detectobjects.py")
+
+    config = procedure.create_config()
+    # config.set_property("force_cpu", force_cpu)
+    # config.begin_run(image, run_mode, args)
+
+    GimpUi.init("filterfolder.py")
+    use_header_bar = Gtk.Settings.get_default().get_property(
+        "gtk-dialogs-use-header"
+    )
+    dialog = GimpUi.Dialog(use_header_bar=use_header_bar, title=_("Filter folder..."))
+    dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
+    dialog.add_button("_Help", Gtk.ResponseType.APPLY)
+    dialog.add_button("_Run Inference", Gtk.ResponseType.OK)
+
+    vbox = Gtk.Box(
+        orientation=Gtk.Orientation.VERTICAL, homogeneous=False, spacing=10
+    )
+    dialog.get_content_area().add(vbox)
+    vbox.show()
+
+    # Create grid to set all the properties inside.
+    grid = Gtk.Grid()
+    grid.set_column_homogeneous(False)
+    grid.set_border_width(10)
+    grid.set_column_spacing(10)
+    grid.set_row_spacing(10)
+    vbox.add(grid)
+    grid.show()
+
+    # # Force CPU parameter
+    # spin = GimpUi.prop_check_button_new(config, "force_cpu", _("Force _CPU"))
+    # spin.set_tooltip_text(
+    #     _(
+    #         "If checked, CPU is used for model inference."
+    #         " Otherwise, GPU will be used if available."
+    #     )
+    # )
+    # grid.attach(spin, 1, 2, 1, 1)
+    # spin.show()
+
+    # UI for the file parameter
+    def choose_file(widget):
+        if file_chooser_dialog.run() == Gtk.ResponseType.OK:
+            if file_chooser_dialog.get_file() is not None:
+                # config.set_property("file", file_chooser_dialog.get_file())
+                file_entry.set_text(file_chooser_dialog.get_file().get_path())
+        file_chooser_dialog.hide()
+
+    file_chooser_button = Gtk.Button.new_with_mnemonic(label=_("_Folder..."))
+    grid.attach(file_chooser_button, 0, 0, 1, 1)
+    file_chooser_button.show()
+    file_chooser_button.connect("clicked", choose_file)
+
+    file_entry = Gtk.Entry.new()
+    grid.attach(file_entry, 1, 0, 1, 1)
+    file_entry.set_width_chars(40)
+    file_entry.set_placeholder_text(_("Choose export folder..."))
+    # if gio_file is not None:
+    file_entry.set_text(os.path.join(os.path.expanduser("~"), "Pictures"))
+    file_entry.show()
+
+    file_chooser_dialog = Gtk.FileChooserDialog(
+        use_header_bar=use_header_bar,
+        title=_("Frame Export folder..."),
+        action=Gtk.FileChooserAction.SELECT_FOLDER,
+    )
+    file_chooser_dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
+    file_chooser_dialog.add_button("_OK", Gtk.ResponseType.OK)
+
+    # Filter Objects
+    filter_objects = Gtk.Entry.new()
+    grid.attach(filter_objects, 1, 1, 1, 1)
+    filter_objects.set_width_chars(40)
+    # file_entry.set_placeholder_text(_("..."))
+    # if gio_file is not None:
+    filter_objects.set_text("Person|Cars")
+    filter_objects.show()
+
+    filer_objects_text = _("Objects to search")
+    filter_label = Gtk.Label(label=filer_objects_text)
+    grid.attach(filter_label, 0, 1, 1, 1)
+    vbox.pack_start(filter_label, False, False, 1)
+    filter_label.show()
+
+    # Show Logo
+    logo = Gtk.Image.new_from_file(image_paths["logo"])
+    # grid.attach(logo, 0, 0, 1, 1)
+    vbox.pack_start(logo, False, False, 1)
+    logo.show()
+
+    # Show Custom Text
+    license_text = _("For complete list of objects see help.")
+    label = Gtk.Label(label=license_text)
+    # grid.attach(label, 1, 2, 1, 1)
+    vbox.pack_start(label, False, False, 1)
+    label.show()
+
+    # Show License
+    license_text = _("PLUGIN LICENSE : Apache-2.0")
+    label = Gtk.Label(label=license_text)
+    # grid.attach(label, 1, 1, 1, 1)
+    vbox.pack_start(label, False, False, 1)
+    label.show()
+
+    progress_bar = Gtk.ProgressBar()
+    vbox.add(progress_bar)
+    progress_bar.show()
+
+    # Wait for user to click
+    dialog.show()
+    while True:
+        response = dialog.run()
+        if response == Gtk.ResponseType.OK:
+            args_dict["image_path"] = file_entry.get_text()
+            args_dict["objects"] = filter_objects.get_text()
+            # force_cpu = config.get_property("force_cpu")
+            result = filterfolder(
+                procedure, args_dict, config_path_output
+            )
+            # super_resolution(procedure, image, n_drawables, layer, force_cpu, progress_bar, config_path_output)
+            # If the execution was successful, save parameters so they will be restored next time we show dialog.
+            # if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
+            #     config.end_run(Gimp.PDBStatusType.SUCCESS)
+            # return result
+            if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
+                config.end_run(Gimp.PDBStatusType.SUCCESS)
+            return result
+        elif response == Gtk.ResponseType.APPLY:
+            url = "https://kritiksoman.github.io/GIMP-ML-Docs/docs-page.html#item-7-15"
+            Gio.app_info_launch_default_for_uri(url, None)
+            continue
+        else:
+            dialog.destroy()
+            return procedure.new_return_values(
+                Gimp.PDBStatusType.CANCEL, GLib.Error()
+            )
+
+
+class FilterFolder(Gimp.PlugIn):
+    ## Properties: parameters ##
+
+    # ## Parameters ##
+    # __gproperties__ = {
+    #     "force_cpu": (
+    #         bool,
+    #         _("Force _CPU"),
+    #         "Force CPU",
+    #         False,
+    #         GObject.ParamFlags.READWRITE,
+    #     ),
+    # }
+
+    @GObject.Property(type=Gimp.RunMode,
+                      default=Gimp.RunMode.NONINTERACTIVE,
+                      nick="Run mode", blurb="The run mode")
+    def run_mode(self):
+        """Read-write integer property."""
+        return self.runmode
+
+    @run_mode.setter
+    def run_mode(self, runmode):
+        self.runmode = runmode
+
+    ## GimpPlugIn virtual methods ##
+    def do_query_procedures(self):
+        self.set_translation_domain(
+            "gimp30-python", Gio.file_new_for_path(Gimp.locale_directory())
+        )
+        return ["filterfolder"]
+
+    def do_create_procedure(self, name):
+        procedure = None
+        if name == "filterfolder":
+            procedure = Gimp.Procedure.new(
+                self, name, Gimp.PDBProcType.PLUGIN, run, None
+            )
+            # procedure.set_image_types("*")
+            procedure.set_documentation(
+                N_("Detects objects on the current layer."),
+                globals()[
+                    "__doc__"
+                ],  # This includes the docstring, on the top of the file
+                name,
+            )
+            procedure.set_menu_label(N_("Filter Folder Objects..."))
+            procedure.set_attribution("Kritik Soman", "GIMP-ML", "2021")
+            procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
+            procedure.add_argument_from_property(self, "run-mode")
+            # procedure.add_argument_from_property(self, "force_cpu")
+
+        return procedure
+
+
+Gimp.main(FilterFolder.__gtype__, sys.argv)

gimpml/tools/canny.py

@@ -0,0 +1,39 @@
+import pickle
+import os
+import cv2
+from gimpml.tools.tools_utils import get_weight_path
+import numpy as np
+
+
+def get_edge(image, min_val=100, max_val=200):
+    if image.shape[2] == 4:  # get rid of alpha channel
+        image = image[:, :, 0:3]
+    edges = cv2.Canny(image, min_val, max_val)
+    edges = np.dstack([edges, edges, edges])
+    return edges
+
+
+if __name__ == "__main__":
+    weight_path = get_weight_path()
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+    min_val = data_output["min_val"]
+    max_val = data_output["max_val"]
+    image = cv2.imread(os.path.join(weight_path, "..", "cache.png"))[:, :, ::-1]
+    try:
+        output = get_edge(image, min_val=min_val, max_val=max_val)
+        cv2.imwrite(os.path.join(weight_path, "..", "cache.png"), output[:, :, ::-1])
+        with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+            pickle.dump({"inference_status": "success"}, file)
+
+        # Remove old temporary error files that were saved
+        my_dir = os.path.join(weight_path, "..")
+        for f_name in os.listdir(my_dir):
+            if f_name.startswith("error_log"):
+                os.remove(os.path.join(my_dir, f_name))
+
+    except Exception as error:
+        with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+            pickle.dump({"inference_status": "failed"}, file)
+        with open(os.path.join(weight_path, "..", "error_log.txt"), "w") as file:
+            file.write(str(error))

gimpml/tools/detectobjects.py

@@ -0,0 +1,134 @@
+import sys
+import os
+
+plugin_loc = os.path.join(os.path.dirname(os.path.realpath(__file__)), "pytorch-YOLOv4")
+sys.path.extend([plugin_loc])
+
+import cv2
+import torch
+from yolo_models import Yolov4
+from yolo_tool.utils import load_class_names, plot_boxes_cv2, get_objects
+from yolo_tool.torch_utils import do_detect
+from gimpml.tools.tools_utils import get_weight_path
+import pickle
+import traceback
+import shutil
+import numpy as np
+
+def scale_image(image):
+    height, width = image.shape[:2]
+    if height > 320 or width > 320:
+        if height > 320:
+            n = (height - 320) / 96
+        if width > 320:
+            m = (width - 320) / 96
+        sized = cv2.resize(image, (320 + 96 * int(np.ceil(m)), 320 + 96 * int(np.ceil(n))))
+    else:
+        sized = cv2.resize(image, (320, 320))
+    scale_width = sized.shape[1] / width
+    scale_height = sized.shape[0] / height
+    return sized, scale_width, scale_height
+
+
+def get_detect_objects(image=None, image_path=None, cpu_flag=False, weight_path=None, get_predict_image=False,
+                   n_classes=80):
+    if image is not None and image_path is not None:
+        raise Exception("Invalid input.")
+
+    if weight_path is None:
+        weight_path = get_weight_path()
+
+    weightfile = os.path.join(weight_path, "yolo", "yolov4.pth")
+
+    use_cuda = False
+    if torch.cuda.is_available() and not cpu_flag:
+        use_cuda = True
+
+    if n_classes == 20:
+        namesfile = os.path.join(os.path.dirname(os.path.realpath(__file__)), "pytorch-YOLOv4", "yolo_data", "voc.names")
+    elif n_classes == 80:
+        namesfile = os.path.join(os.path.dirname(os.path.realpath(__file__)), "pytorch-YOLOv4", "yolo_data", "coco.names")
+    else:
+        print("please give namefile")
+    class_names = load_class_names(namesfile)
+
+    model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True)
+    if use_cuda:
+        pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
+    else:
+        pretrained_dict = torch.load(weightfile, map_location=torch.device('cpu'))
+    model.load_state_dict(pretrained_dict)
+    if use_cuda:
+        model.cuda()
+
+    result = []
+    if image is not None and image_path is None:
+        sized, scale_width, scale_height = scale_image(image)
+        with torch.no_grad():
+            boxes = do_detect(model, sized, 0.4, 0.6, use_cuda)
+        result.append(plot_boxes_cv2(image, boxes[0], scale_width, scale_height, class_names=class_names) if get_predict_image else
+                      (get_objects(image, boxes[0], scale_width, scale_height, class_names), "image"))
+    elif image is None and image_path is not None:
+        for filename in os.listdir(image_path):
+            try:
+                image = cv2.imread(os.path.join(image_path, filename))[:, :, [2, 1, 0]]
+                sized, scale_width, scale_height = scale_image(image)
+                with torch.no_grad():
+                    boxes = do_detect(model, sized, 0.4, 0.6, use_cuda)
+                result.append(plot_boxes_cv2(sized, boxes[0], scale_width, scale_height, class_names=class_names) if get_predict_image else
+                              (get_objects(sized, boxes[0], scale_width, scale_height, class_names),
+                               os.path.join(image_path, filename)))
+            except:
+                pass
+    return result
+
+
+if __name__ == "__main__":
+    weight_path = get_weight_path()
+    with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "rb") as file:
+        data_output = pickle.load(file)
+    force_cpu = False #data_output["force_cpu"]
+    get_predict_image = data_output["get_predict_image"]
+    image1, image_path = None, None
+    if get_predict_image:
+        image1 = cv2.imread(os.path.join(weight_path, "..", "cache.png"))
+    else:
+        image_path = data_output["image_path"]
+        search_objects = [x.lower().strip() for x in data_output["objects"].split("|")]
+    try:
+        if get_predict_image:
+            count = 0
+            output = get_detect_objects(image=image1, cpu_flag=force_cpu, weight_path=weight_path, get_predict_image=True)[0]
+            # with open(os.path.join(weight_path, "..", "output.txt"), "w") as file:
+            #     file.write(str(output))
+            cv2.imwrite(os.path.join(weight_path, "..", "cache.png"), output)
+        else:
+            count = 0
+            output = get_detect_objects(image_path=image_path, cpu_flag=force_cpu, weight_path=weight_path)
+            save_filtered_path = os.path.join(image_path, "filtered")
+            if not os.path.exists(save_filtered_path):
+                os.makedirs(save_filtered_path)
+            for res in output:
+                if any([obj[-2] in search_objects for obj in res[0]]):
+                    head, tail = os.path.split(res[-1])
+                    shutil.move(res[-1], os.path.join(head, "filtered", tail))
+                    count += 1
+            # with open(os.path.join(weight_path, "..", "tmp.txt"), "w") as file:
+            #     file.write(str(output))
+
+        with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+            pickle.dump({"inference_status": "success", "force_cpu": force_cpu, "count": count}, file)
+
+        # Remove old temporary error files that were saved
+        my_dir = os.path.join(weight_path, "..")
+        for f_name in os.listdir(my_dir):
+            if f_name.startswith("error_log"):
+                os.remove(os.path.join(my_dir, f_name))
+
+    except Exception as error:
+        with open(os.path.join(weight_path, "..", "gimp_ml_run.pkl"), "wb") as file:
+            pickle.dump({"inference_status": "failed"}, file)
+        with open(os.path.join(weight_path, "..", "error_log.txt"), "w") as file:
+            # file.write(str(error))
+            e_type, e_val, e_tb = sys.exc_info()
+            traceback.print_exception(e_type, e_val, e_tb, file=file)

gimpml/tools/model_info.csv

@@ -26,3 +26,4 @@ edgeconnect/psv,11vIrpFaAHhgPN0yixFbN7ZNY78ooefiq,10.8,EdgeModel_dis.pth,ea6c734
 edgeconnect/psv,11taSeBwJXj3AC1qucNHFUgTOTDIbyOLz,42,EdgeModel_gen.pth,40c83e91ace061acd2bd95421750ffaf
 edgeconnect/psv,122BaS0KobaZqsU5mZcIR4pTsoQz__1k8,10.8,InpaintingModel_dis.pth,bf9bb863592605237e620b8d73db225e
 edgeconnect/psv,11wPatO4UAIuPc_9Se99QvlKwGRiw62PY,42,InpaintingModel_gen.pth,afc0ce9b90413298972a2ef1fc65a3c7
+yolo,1-1m-rpIsbolCUGZRdsrKXq-KsT4A-Iwl,246,yolov4.pth,fa40d42ad9f3eb35abdee056aef3f936

gimpml/tools/pytorch-YOLOv4/License.txt

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

gimpml/tools/pytorch-YOLOv4/yolo_data/coco.names

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorbike
+aeroplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+sofa
+pottedplant
+bed
+diningtable
+toilet
+tvmonitor
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

gimpml/tools/pytorch-YOLOv4/yolo_data/voc.names

@@ -0,0 +1,20 @@
+aeroplane
+bicycle
+bird
+boat
+bottle
+bus
+car
+cat
+chair
+cow
+diningtable
+dog
+horse
+motorbike
+person
+pottedplant
+sheep
+sofa
+train
+tvmonitor

gimpml/tools/pytorch-YOLOv4/yolo_models.py

@@ -0,0 +1,450 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from yolo_tool.torch_utils import *
+from yolo_tool.yolo_layer import YoloLayer
+
+
+class Mish(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x = x * (torch.tanh(torch.nn.functional.softplus(x)))
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self):
+        super(Upsample, self).__init__()
+
+    def forward(self, x, target_size, inference=False):
+        assert (x.data.dim() == 4)
+        # _, _, tH, tW = target_size
+
+        if inference:
+
+            #B = x.data.size(0)
+            #C = x.data.size(1)
+            #H = x.data.size(2)
+            #W = x.data.size(3)
+
+            return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
+                    expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3), target_size[3] // x.size(3)).\
+                    contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
+        else:
+            return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')
+
+
+class Conv_Bn_Activation(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, activation, bn=True, bias=False):
+        super().__init__()
+        pad = (kernel_size - 1) // 2
+
+        self.conv = nn.ModuleList()
+        if bias:
+            self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad))
+        else:
+            self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad, bias=False))
+        if bn:
+            self.conv.append(nn.BatchNorm2d(out_channels))
+        if activation == "mish":
+            self.conv.append(Mish())
+        elif activation == "relu":
+            self.conv.append(nn.ReLU(inplace=True))
+        elif activation == "leaky":
+            self.conv.append(nn.LeakyReLU(0.1, inplace=True))
+        elif activation == "linear":
+            pass
+        else:
+            print("activate error !!! {} {} {}".format(sys._getframe().f_code.co_filename,
+                                                       sys._getframe().f_code.co_name, sys._getframe().f_lineno))
+
+    def forward(self, x):
+        for l in self.conv:
+            x = l(x)
+        return x
+
+
+class ResBlock(nn.Module):
+    """
+    Sequential residual blocks each of which consists of \
+    two convolution layers.
+    Args:
+        ch (int): number of input and output channels.
+        nblocks (int): number of residual blocks.
+        shortcut (bool): if True, residual tensor addition is enabled.
+    """
+
+    def __init__(self, ch, nblocks=1, shortcut=True):
+        super().__init__()
+        self.shortcut = shortcut
+        self.module_list = nn.ModuleList()
+        for i in range(nblocks):
+            resblock_one = nn.ModuleList()
+            resblock_one.append(Conv_Bn_Activation(ch, ch, 1, 1, 'mish'))
+            resblock_one.append(Conv_Bn_Activation(ch, ch, 3, 1, 'mish'))
+            self.module_list.append(resblock_one)
+
+    def forward(self, x):
+        for module in self.module_list:
+            h = x
+            for res in module:
+                h = res(h)
+            x = x + h if self.shortcut else h
+        return x
+
+
+class DownSample1(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = Conv_Bn_Activation(3, 32, 3, 1, 'mish')
+
+        self.conv2 = Conv_Bn_Activation(32, 64, 3, 2, 'mish')
+        self.conv3 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
+        # [route]
+        # layers = -2
+        self.conv4 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
+
+        self.conv5 = Conv_Bn_Activation(64, 32, 1, 1, 'mish')
+        self.conv6 = Conv_Bn_Activation(32, 64, 3, 1, 'mish')
+        # [shortcut]
+        # from=-3
+        # activation = linear
+
+        self.conv7 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
+        # [route]
+        # layers = -1, -7
+        self.conv8 = Conv_Bn_Activation(128, 64, 1, 1, 'mish')
+
+    def forward(self, input):
+        x1 = self.conv1(input)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x2)
+        # route -2
+        x4 = self.conv4(x2)
+        x5 = self.conv5(x4)
+        x6 = self.conv6(x5)
+        # shortcut -3
+        x6 = x6 + x4
+
+        x7 = self.conv7(x6)
+        # [route]
+        # layers = -1, -7
+        x7 = torch.cat([x7, x3], dim=1)
+        x8 = self.conv8(x7)
+        return x8
+
+
+class DownSample2(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = Conv_Bn_Activation(64, 128, 3, 2, 'mish')
+        self.conv2 = Conv_Bn_Activation(128, 64, 1, 1, 'mish')
+        # r -2
+        self.conv3 = Conv_Bn_Activation(128, 64, 1, 1, 'mish')
+
+        self.resblock = ResBlock(ch=64, nblocks=2)
+
+        # s -3
+        self.conv4 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
+        # r -1 -10
+        self.conv5 = Conv_Bn_Activation(128, 128, 1, 1, 'mish')
+
+    def forward(self, input):
+        x1 = self.conv1(input)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x1)
+
+        r = self.resblock(x3)
+        x4 = self.conv4(r)
+
+        x4 = torch.cat([x4, x2], dim=1)
+        x5 = self.conv5(x4)
+        return x5
+
+
+class DownSample3(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = Conv_Bn_Activation(128, 256, 3, 2, 'mish')
+        self.conv2 = Conv_Bn_Activation(256, 128, 1, 1, 'mish')
+        self.conv3 = Conv_Bn_Activation(256, 128, 1, 1, 'mish')
+
+        self.resblock = ResBlock(ch=128, nblocks=8)
+        self.conv4 = Conv_Bn_Activation(128, 128, 1, 1, 'mish')
+        self.conv5 = Conv_Bn_Activation(256, 256, 1, 1, 'mish')
+
+    def forward(self, input):
+        x1 = self.conv1(input)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x1)
+
+        r = self.resblock(x3)
+        x4 = self.conv4(r)
+
+        x4 = torch.cat([x4, x2], dim=1)
+        x5 = self.conv5(x4)
+        return x5
+
+
+class DownSample4(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = Conv_Bn_Activation(256, 512, 3, 2, 'mish')
+        self.conv2 = Conv_Bn_Activation(512, 256, 1, 1, 'mish')
+        self.conv3 = Conv_Bn_Activation(512, 256, 1, 1, 'mish')
+
+        self.resblock = ResBlock(ch=256, nblocks=8)
+        self.conv4 = Conv_Bn_Activation(256, 256, 1, 1, 'mish')
+        self.conv5 = Conv_Bn_Activation(512, 512, 1, 1, 'mish')
+
+    def forward(self, input):
+        x1 = self.conv1(input)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x1)
+
+        r = self.resblock(x3)
+        x4 = self.conv4(r)
+
+        x4 = torch.cat([x4, x2], dim=1)
+        x5 = self.conv5(x4)
+        return x5
+
+
+class DownSample5(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = Conv_Bn_Activation(512, 1024, 3, 2, 'mish')
+        self.conv2 = Conv_Bn_Activation(1024, 512, 1, 1, 'mish')
+        self.conv3 = Conv_Bn_Activation(1024, 512, 1, 1, 'mish')
+
+        self.resblock = ResBlock(ch=512, nblocks=4)
+        self.conv4 = Conv_Bn_Activation(512, 512, 1, 1, 'mish')
+        self.conv5 = Conv_Bn_Activation(1024, 1024, 1, 1, 'mish')
+
+    def forward(self, input):
+        x1 = self.conv1(input)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x1)
+
+        r = self.resblock(x3)
+        x4 = self.conv4(r)
+
+        x4 = torch.cat([x4, x2], dim=1)
+        x5 = self.conv5(x4)
+        return x5
+
+
+class Neck(nn.Module):
+    def __init__(self, inference=False):
+        super().__init__()
+        self.inference = inference
+
+        self.conv1 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
+        self.conv2 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
+        self.conv3 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
+        # SPP
+        self.maxpool1 = nn.MaxPool2d(kernel_size=5, stride=1, padding=5 // 2)
+        self.maxpool2 = nn.MaxPool2d(kernel_size=9, stride=1, padding=9 // 2)
+        self.maxpool3 = nn.MaxPool2d(kernel_size=13, stride=1, padding=13 // 2)
+
+        # R -1 -3 -5 -6
+        # SPP
+        self.conv4 = Conv_Bn_Activation(2048, 512, 1, 1, 'leaky')
+        self.conv5 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
+        self.conv6 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
+        self.conv7 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        # UP
+        self.upsample1 = Upsample()
+        # R 85
+        self.conv8 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        # R -1 -3
+        self.conv9 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        self.conv10 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
+        self.conv11 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        self.conv12 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
+        self.conv13 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        self.conv14 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
+        # UP
+        self.upsample2 = Upsample()
+        # R 54
+        self.conv15 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
+        # R -1 -3
+        self.conv16 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
+        self.conv17 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
+        self.conv18 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
+        self.conv19 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
+        self.conv20 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
+
+    def forward(self, input, downsample4, downsample3, inference=False):
+        x1 = self.conv1(input)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x2)
+        # SPP
+        m1 = self.maxpool1(x3)
+        m2 = self.maxpool2(x3)
+        m3 = self.maxpool3(x3)
+        spp = torch.cat([m3, m2, m1, x3], dim=1)
+        # SPP end
+        x4 = self.conv4(spp)
+        x5 = self.conv5(x4)
+        x6 = self.conv6(x5)
+        x7 = self.conv7(x6)
+        # UP
+        up = self.upsample1(x7, downsample4.size(), self.inference)
+        # R 85
+        x8 = self.conv8(downsample4)
+        # R -1 -3
+        x8 = torch.cat([x8, up], dim=1)
+
+        x9 = self.conv9(x8)
+        x10 = self.conv10(x9)
+        x11 = self.conv11(x10)
+        x12 = self.conv12(x11)
+        x13 = self.conv13(x12)
+        x14 = self.conv14(x13)
+
+        # UP
+        up = self.upsample2(x14, downsample3.size(), self.inference)
+        # R 54
+        x15 = self.conv15(downsample3)
+        # R -1 -3
+        x15 = torch.cat([x15, up], dim=1)
+
+        x16 = self.conv16(x15)
+        x17 = self.conv17(x16)
+        x18 = self.conv18(x17)
+        x19 = self.conv19(x18)
+        x20 = self.conv20(x19)
+        return x20, x13, x6
+
+
+class Yolov4Head(nn.Module):
+    def __init__(self, output_ch, n_classes, inference=False):
+        super().__init__()
+        self.inference = inference
+
+        self.conv1 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
+        self.conv2 = Conv_Bn_Activation(256, output_ch, 1, 1, 'linear', bn=False, bias=True)
+
+        self.yolo1 = YoloLayer(
+                                anchor_mask=[0, 1, 2], num_classes=n_classes,
+                                anchors=[12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401],
+                                num_anchors=9, stride=8)
+
+        # R -4
+        self.conv3 = Conv_Bn_Activation(128, 256, 3, 2, 'leaky')
+
+        # R -1 -16
+        self.conv4 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        self.conv5 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
+        self.conv6 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        self.conv7 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
+        self.conv8 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
+        self.conv9 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
+        self.conv10 = Conv_Bn_Activation(512, output_ch, 1, 1, 'linear', bn=False, bias=True)
+        
+        self.yolo2 = YoloLayer(
+                                anchor_mask=[3, 4, 5], num_classes=n_classes,
+                                anchors=[12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401],
+                                num_anchors=9, stride=16)
+
+        # R -4
+        self.conv11 = Conv_Bn_Activation(256, 512, 3, 2, 'leaky')
+
+        # R -1 -37
+        self.conv12 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
+        self.conv13 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
+        self.conv14 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
+        self.conv15 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
+        self.conv16 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
+        self.conv17 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
+        self.conv18 = Conv_Bn_Activation(1024, output_ch, 1, 1, 'linear', bn=False, bias=True)
+        
+        self.yolo3 = YoloLayer(
+                                anchor_mask=[6, 7, 8], num_classes=n_classes,
+                                anchors=[12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401],
+                                num_anchors=9, stride=32)
+
+    def forward(self, input1, input2, input3):
+        x1 = self.conv1(input1)
+        x2 = self.conv2(x1)
+
+        x3 = self.conv3(input1)
+        # R -1 -16
+        x3 = torch.cat([x3, input2], dim=1)
+        x4 = self.conv4(x3)
+        x5 = self.conv5(x4)
+        x6 = self.conv6(x5)
+        x7 = self.conv7(x6)
+        x8 = self.conv8(x7)
+        x9 = self.conv9(x8)
+        x10 = self.conv10(x9)
+
+        # R -4
+        x11 = self.conv11(x8)
+        # R -1 -37
+        x11 = torch.cat([x11, input3], dim=1)
+
+        x12 = self.conv12(x11)
+        x13 = self.conv13(x12)
+        x14 = self.conv14(x13)
+        x15 = self.conv15(x14)
+        x16 = self.conv16(x15)
+        x17 = self.conv17(x16)
+        x18 = self.conv18(x17)
+        
+        if self.inference:
+            y1 = self.yolo1(x2)
+            y2 = self.yolo2(x10)
+            y3 = self.yolo3(x18)
+
+            return get_region_boxes([y1, y2, y3])
+        
+        else:
+            return [x2, x10, x18]
+
+
+class Yolov4(nn.Module):
+    def __init__(self, yolov4conv137weight=None, n_classes=80, inference=False):
+        super().__init__()
+
+        output_ch = (4 + 1 + n_classes) * 3
+
+        # backbone
+        self.down1 = DownSample1()
+        self.down2 = DownSample2()
+        self.down3 = DownSample3()
+        self.down4 = DownSample4()
+        self.down5 = DownSample5()
+        # neck
+        self.neek = Neck(inference)
+        # yolov4conv137
+        if yolov4conv137weight:
+            _model = nn.Sequential(self.down1, self.down2, self.down3, self.down4, self.down5, self.neek)
+            pretrained_dict = torch.load(yolov4conv137weight)
+
+            model_dict = _model.state_dict()
+            # 1. filter out unnecessary keys
+            pretrained_dict = {k1: v for (k, v), k1 in zip(pretrained_dict.items(), model_dict)}
+            # 2. overwrite entries in the existing state dict
+            model_dict.update(pretrained_dict)
+            _model.load_state_dict(model_dict)
+        
+        # head
+        self.head = Yolov4Head(output_ch, n_classes, inference)
+
+
+    def forward(self, input):
+        d1 = self.down1(input)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+        d5 = self.down5(d4)
+
+        x20, x13, x6 = self.neek(d5, d4, d3)
+
+        output = self.head(x20, x13, x6)
+        return output
+

gimpml/tools/pytorch-YOLOv4/yolo_tool/torch_utils.py

@@ -0,0 +1,104 @@
+import sys
+import os
+import time
+import math
+import torch
+import numpy as np
+from torch.autograd import Variable
+
+import itertools
+import struct  # get_image_size
+import imghdr  # get_image_size
+
+from yolo_tool import utils
+
+
+def bbox_ious(boxes1, boxes2, x1y1x2y2=True):
+    if x1y1x2y2:
+        mx = torch.min(boxes1[0], boxes2[0])
+        Mx = torch.max(boxes1[2], boxes2[2])
+        my = torch.min(boxes1[1], boxes2[1])
+        My = torch.max(boxes1[3], boxes2[3])
+        w1 = boxes1[2] - boxes1[0]
+        h1 = boxes1[3] - boxes1[1]
+        w2 = boxes2[2] - boxes2[0]
+        h2 = boxes2[3] - boxes2[1]
+    else:
+        mx = torch.min(boxes1[0] - boxes1[2] / 2.0, boxes2[0] - boxes2[2] / 2.0)
+        Mx = torch.max(boxes1[0] + boxes1[2] / 2.0, boxes2[0] + boxes2[2] / 2.0)
+        my = torch.min(boxes1[1] - boxes1[3] / 2.0, boxes2[1] - boxes2[3] / 2.0)
+        My = torch.max(boxes1[1] + boxes1[3] / 2.0, boxes2[1] + boxes2[3] / 2.0)
+        w1 = boxes1[2]
+        h1 = boxes1[3]
+        w2 = boxes2[2]
+        h2 = boxes2[3]
+    uw = Mx - mx
+    uh = My - my
+    cw = w1 + w2 - uw
+    ch = h1 + h2 - uh
+    mask = ((cw <= 0) + (ch <= 0) > 0)
+    area1 = w1 * h1
+    area2 = w2 * h2
+    carea = cw * ch
+    carea[mask] = 0
+    uarea = area1 + area2 - carea
+    return carea / uarea
+
+
+def get_region_boxes(boxes_and_confs):
+
+    # print('Getting boxes from boxes and confs ...')
+
+    boxes_list = []
+    confs_list = []
+
+    for item in boxes_and_confs:
+        boxes_list.append(item[0])
+        confs_list.append(item[1])
+
+    # boxes: [batch, num1 + num2 + num3, 1, 4]
+    # confs: [batch, num1 + num2 + num3, num_classes]
+    boxes = torch.cat(boxes_list, dim=1)
+    confs = torch.cat(confs_list, dim=1)
+        
+    return [boxes, confs]
+
+
+def convert2cpu(gpu_matrix):
+    return torch.FloatTensor(gpu_matrix.size()).copy_(gpu_matrix)
+
+
+def convert2cpu_long(gpu_matrix):
+    return torch.LongTensor(gpu_matrix.size()).copy_(gpu_matrix)
+
+
+
+def do_detect(model, img, conf_thresh, nms_thresh, use_cuda=1):
+    model.eval()
+    t0 = time.time()
+
+    if type(img) == np.ndarray and len(img.shape) == 3:  # cv2 image
+        img = torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0)
+    elif type(img) == np.ndarray and len(img.shape) == 4:
+        img = torch.from_numpy(img.transpose(0, 3, 1, 2)).float().div(255.0)
+    else:
+        print("unknow image type")
+        exit(-1)
+
+    if use_cuda:
+        img = img.cuda()
+    img = torch.autograd.Variable(img)
+    
+    t1 = time.time()
+
+    output = model(img)
+
+    t2 = time.time()
+    #
+    # print('-----------------------------------')
+    # print('           Preprocess : %f' % (t1 - t0))
+    # print('      Model Inference : %f' % (t2 - t1))
+    # print('-----------------------------------')
+
+    return utils.post_processing(img, conf_thresh, nms_thresh, output)
+

gimpml/tools/pytorch-YOLOv4/yolo_tool/utils.py

@@ -0,0 +1,248 @@
+import sys
+import os
+import time
+import math
+import numpy as np
+
+import itertools
+import struct  # get_image_size
+import imghdr  # get_image_size
+
+
+def sigmoid(x):
+    return 1.0 / (np.exp(-x) + 1.)
+
+
+def softmax(x):
+    x = np.exp(x - np.expand_dims(np.max(x, axis=1), axis=1))
+    x = x / np.expand_dims(x.sum(axis=1), axis=1)
+    return x
+
+
+def bbox_iou(box1, box2, x1y1x2y2=True):
+    # print('iou box1:', box1)
+    # print('iou box2:', box2)
+
+    if x1y1x2y2:
+        mx = min(box1[0], box2[0])
+        Mx = max(box1[2], box2[2])
+        my = min(box1[1], box2[1])
+        My = max(box1[3], box2[3])
+        w1 = box1[2] - box1[0]
+        h1 = box1[3] - box1[1]
+        w2 = box2[2] - box2[0]
+        h2 = box2[3] - box2[1]
+    else:
+        w1 = box1[2]
+        h1 = box1[3]
+        w2 = box2[2]
+        h2 = box2[3]
+
+        mx = min(box1[0], box2[0])
+        Mx = max(box1[0] + w1, box2[0] + w2)
+        my = min(box1[1], box2[1])
+        My = max(box1[1] + h1, box2[1] + h2)
+    uw = Mx - mx
+    uh = My - my
+    cw = w1 + w2 - uw
+    ch = h1 + h2 - uh
+    carea = 0
+    if cw <= 0 or ch <= 0:
+        return 0.0
+
+    area1 = w1 * h1
+    area2 = w2 * h2
+    carea = cw * ch
+    uarea = area1 + area2 - carea
+    return carea / uarea
+
+
+def nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False):
+    # print(boxes.shape)
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1) * (y2 - y1)
+    order = confs.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        idx_self = order[0]
+        idx_other = order[1:]
+
+        keep.append(idx_self)
+
+        xx1 = np.maximum(x1[idx_self], x1[idx_other])
+        yy1 = np.maximum(y1[idx_self], y1[idx_other])
+        xx2 = np.minimum(x2[idx_self], x2[idx_other])
+        yy2 = np.minimum(y2[idx_self], y2[idx_other])
+
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+
+        if min_mode:
+            over = inter / np.minimum(areas[order[0]], areas[order[1:]])
+        else:
+            over = inter / (areas[order[0]] + areas[order[1:]] - inter)
+
+        inds = np.where(over <= nms_thresh)[0]
+        order = order[inds + 1]
+
+    return np.array(keep)
+
+
+def plot_boxes_cv2(img, boxes, scale_width, scale_height, savename=None, class_names=None, color=None):
+    import cv2
+    img = np.copy(img)
+    colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)
+
+    def get_color(c, x, max_val):
+        ratio = float(x) / max_val * 5
+        i = int(math.floor(ratio))
+        j = int(math.ceil(ratio))
+        ratio = ratio - i
+        r = (1 - ratio) * colors[i][c] + ratio * colors[j][c]
+        return int(r * 255)
+
+    width = img.shape[1]
+    height = img.shape[0]
+    for i in range(len(boxes)):
+        box = boxes[i]
+        x1 = int(box[0] * width * scale_width)
+        y1 = int(box[1] * height * scale_height)
+        x2 = int(box[2] * width * scale_width)
+        y2 = int(box[3] * height * scale_height)
+
+        if color:
+            rgb = color
+        else:
+            rgb = (255, 0, 0)
+        if len(box) >= 7 and class_names:
+            cls_conf = box[5]
+            cls_id = box[6]
+            # print('%s: %f' % (class_names[cls_id], cls_conf))
+            classes = len(class_names)
+            offset = cls_id * 123457 % classes
+            red = get_color(2, offset, classes)
+            green = get_color(1, offset, classes)
+            blue = get_color(0, offset, classes)
+            if color is None:
+                rgb = (red, green, blue)
+            img = cv2.putText(img, class_names[cls_id], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1.2, rgb, 1)
+        img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, 1)
+    if savename:
+        print("save plot results to %s" % savename)
+        cv2.imwrite(savename, img)
+    return img
+
+
+def get_objects(img, boxes, scale_width, scale_height, class_names=None):
+    objects = []
+    width = img.shape[1]
+    height = img.shape[0]
+    for i in range(len(boxes)):
+        box = boxes[i]
+        x1 = int(box[0] * width * scale_width)
+        y1 = int(box[1] * height * scale_height)
+        x2 = int(box[2] * width * scale_width)
+        y2 = int(box[3] * height * scale_height)
+        cls_conf = box[5]
+        cls_id = box[6]
+        cls = class_names[cls_id]
+        objects.append((x1, y1, x2, y2, cls, cls_conf))
+    return objects
+
+
+def read_truths(lab_path):
+    if not os.path.exists(lab_path):
+        return np.array([])
+    if os.path.getsize(lab_path):
+        truths = np.loadtxt(lab_path)
+        truths = truths.reshape(truths.size / 5, 5)  # to avoid single truth problem
+        return truths
+    else:
+        return np.array([])
+
+
+def load_class_names(namesfile):
+    class_names = []
+    with open(namesfile, 'r') as fp:
+        lines = fp.readlines()
+    for line in lines:
+        line = line.rstrip()
+        class_names.append(line)
+    return class_names
+
+
+def post_processing(img, conf_thresh, nms_thresh, output):
+    # anchors = [12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401]
+    # num_anchors = 9
+    # anchor_masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+    # strides = [8, 16, 32]
+    # anchor_step = len(anchors) // num_anchors
+
+    # [batch, num, 1, 4]
+    box_array = output[0]
+    # [batch, num, num_classes]
+    confs = output[1]
+
+    t1 = time.time()
+
+    if type(box_array).__name__ != 'ndarray':
+        box_array = box_array.cpu().detach().numpy()
+        confs = confs.cpu().detach().numpy()
+
+    num_classes = confs.shape[2]
+
+    # [batch, num, 4]
+    box_array = box_array[:, :, 0]
+
+    # [batch, num, num_classes] --> [batch, num]
+    max_conf = np.max(confs, axis=2)
+    max_id = np.argmax(confs, axis=2)
+
+    t2 = time.time()
+
+    bboxes_batch = []
+    for i in range(box_array.shape[0]):
+
+        argwhere = max_conf[i] > conf_thresh
+        l_box_array = box_array[i, argwhere, :]
+        l_max_conf = max_conf[i, argwhere]
+        l_max_id = max_id[i, argwhere]
+
+        bboxes = []
+        # nms for each class
+        for j in range(num_classes):
+
+            cls_argwhere = l_max_id == j
+            ll_box_array = l_box_array[cls_argwhere, :]
+            ll_max_conf = l_max_conf[cls_argwhere]
+            ll_max_id = l_max_id[cls_argwhere]
+
+            keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)
+
+            if (keep.size > 0):
+                ll_box_array = ll_box_array[keep, :]
+                ll_max_conf = ll_max_conf[keep]
+                ll_max_id = ll_max_id[keep]
+
+                for k in range(ll_box_array.shape[0]):
+                    bboxes.append(
+                        [ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k],
+                         ll_max_conf[k], ll_max_id[k]])
+
+        bboxes_batch.append(bboxes)
+
+    t3 = time.time()
+
+    # print('-----------------------------------')
+    # print('       max and argmax : %f' % (t2 - t1))
+    # print('                  nms : %f' % (t3 - t2))
+    # print('Post processing total : %f' % (t3 - t1))
+    # print('-----------------------------------')
+
+    return bboxes_batch

gimpml/tools/pytorch-YOLOv4/yolo_tool/yolo_layer.py

@@ -0,0 +1,322 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from yolo_tool.torch_utils import *
+
+def yolo_forward(output, conf_thresh, num_classes, anchors, num_anchors, scale_x_y, only_objectness=1,
+                              validation=False):
+    # Output would be invalid if it does not satisfy this assert
+    # assert (output.size(1) == (5 + num_classes) * num_anchors)
+
+    # print(output.size())
+
+    # Slice the second dimension (channel) of output into:
+    # [ 2, 2, 1, num_classes, 2, 2, 1, num_classes, 2, 2, 1, num_classes ]
+    # And then into
+    # bxy = [ 6 ] bwh = [ 6 ] det_conf = [ 3 ] cls_conf = [ num_classes * 3 ]
+    batch = output.size(0)
+    H = output.size(2)
+    W = output.size(3)
+
+    bxy_list = []
+    bwh_list = []
+    det_confs_list = []
+    cls_confs_list = []
+
+    for i in range(num_anchors):
+        begin = i * (5 + num_classes)
+        end = (i + 1) * (5 + num_classes)
+        
+        bxy_list.append(output[:, begin : begin + 2])
+        bwh_list.append(output[:, begin + 2 : begin + 4])
+        det_confs_list.append(output[:, begin + 4 : begin + 5])
+        cls_confs_list.append(output[:, begin + 5 : end])
+
+    # Shape: [batch, num_anchors * 2, H, W]
+    bxy = torch.cat(bxy_list, dim=1)
+    # Shape: [batch, num_anchors * 2, H, W]
+    bwh = torch.cat(bwh_list, dim=1)
+
+    # Shape: [batch, num_anchors, H, W]
+    det_confs = torch.cat(det_confs_list, dim=1)
+    # Shape: [batch, num_anchors * H * W]
+    det_confs = det_confs.view(batch, num_anchors * H * W)
+
+    # Shape: [batch, num_anchors * num_classes, H, W]
+    cls_confs = torch.cat(cls_confs_list, dim=1)
+    # Shape: [batch, num_anchors, num_classes, H * W]
+    cls_confs = cls_confs.view(batch, num_anchors, num_classes, H * W)
+    # Shape: [batch, num_anchors, num_classes, H * W] --> [batch, num_anchors * H * W, num_classes] 
+    cls_confs = cls_confs.permute(0, 1, 3, 2).reshape(batch, num_anchors * H * W, num_classes)
+
+    # Apply sigmoid(), exp() and softmax() to slices
+    #
+    bxy = torch.sigmoid(bxy) * scale_x_y - 0.5 * (scale_x_y - 1)
+    bwh = torch.exp(bwh)
+    det_confs = torch.sigmoid(det_confs)
+    cls_confs = torch.sigmoid(cls_confs)
+
+    # Prepare C-x, C-y, P-w, P-h (None of them are torch related)
+    grid_x = np.expand_dims(np.expand_dims(np.expand_dims(np.linspace(0, W - 1, W), axis=0).repeat(H, 0), axis=0), axis=0)
+    grid_y = np.expand_dims(np.expand_dims(np.expand_dims(np.linspace(0, H - 1, H), axis=1).repeat(W, 1), axis=0), axis=0)
+    # grid_x = torch.linspace(0, W - 1, W).reshape(1, 1, 1, W).repeat(1, 1, H, 1)
+    # grid_y = torch.linspace(0, H - 1, H).reshape(1, 1, H, 1).repeat(1, 1, 1, W)
+
+    anchor_w = []
+    anchor_h = []
+    for i in range(num_anchors):
+        anchor_w.append(anchors[i * 2])
+        anchor_h.append(anchors[i * 2 + 1])
+
+    device = None
+    cuda_check = output.is_cuda
+    if cuda_check:
+        device = output.get_device()
+
+    bx_list = []
+    by_list = []
+    bw_list = []
+    bh_list = []
+
+    # Apply C-x, C-y, P-w, P-h
+    for i in range(num_anchors):
+        ii = i * 2
+        # Shape: [batch, 1, H, W]
+        bx = bxy[:, ii : ii + 1] + torch.tensor(grid_x, device=device, dtype=torch.float32) # grid_x.to(device=device, dtype=torch.float32)
+        # Shape: [batch, 1, H, W]
+        by = bxy[:, ii + 1 : ii + 2] + torch.tensor(grid_y, device=device, dtype=torch.float32) # grid_y.to(device=device, dtype=torch.float32)
+        # Shape: [batch, 1, H, W]
+        bw = bwh[:, ii : ii + 1] * anchor_w[i]
+        # Shape: [batch, 1, H, W]
+        bh = bwh[:, ii + 1 : ii + 2] * anchor_h[i]
+
+        bx_list.append(bx)
+        by_list.append(by)
+        bw_list.append(bw)
+        bh_list.append(bh)
+
+
+    ########################################
+    #   Figure out bboxes from slices     #
+    ########################################
+    
+    # Shape: [batch, num_anchors, H, W]
+    bx = torch.cat(bx_list, dim=1)
+    # Shape: [batch, num_anchors, H, W]
+    by = torch.cat(by_list, dim=1)
+    # Shape: [batch, num_anchors, H, W]
+    bw = torch.cat(bw_list, dim=1)
+    # Shape: [batch, num_anchors, H, W]
+    bh = torch.cat(bh_list, dim=1)
+
+    # Shape: [batch, 2 * num_anchors, H, W]
+    bx_bw = torch.cat((bx, bw), dim=1)
+    # Shape: [batch, 2 * num_anchors, H, W]
+    by_bh = torch.cat((by, bh), dim=1)
+
+    # normalize coordinates to [0, 1]
+    bx_bw /= W
+    by_bh /= H
+
+    # Shape: [batch, num_anchors * H * W, 1]
+    bx = bx_bw[:, :num_anchors].view(batch, num_anchors * H * W, 1)
+    by = by_bh[:, :num_anchors].view(batch, num_anchors * H * W, 1)
+    bw = bx_bw[:, num_anchors:].view(batch, num_anchors * H * W, 1)
+    bh = by_bh[:, num_anchors:].view(batch, num_anchors * H * W, 1)
+
+    bx1 = bx - bw * 0.5
+    by1 = by - bh * 0.5
+    bx2 = bx1 + bw
+    by2 = by1 + bh
+
+    # Shape: [batch, num_anchors * h * w, 4] -> [batch, num_anchors * h * w, 1, 4]
+    boxes = torch.cat((bx1, by1, bx2, by2), dim=2).view(batch, num_anchors * H * W, 1, 4)
+    # boxes = boxes.repeat(1, 1, num_classes, 1)
+
+    # boxes:     [batch, num_anchors * H * W, 1, 4]
+    # cls_confs: [batch, num_anchors * H * W, num_classes]
+    # det_confs: [batch, num_anchors * H * W]
+
+    det_confs = det_confs.view(batch, num_anchors * H * W, 1)
+    confs = cls_confs * det_confs
+
+    # boxes: [batch, num_anchors * H * W, 1, 4]
+    # confs: [batch, num_anchors * H * W, num_classes]
+
+    return  boxes, confs
+
+
+def yolo_forward_dynamic(output, conf_thresh, num_classes, anchors, num_anchors, scale_x_y, only_objectness=1,
+                              validation=False):
+    # Output would be invalid if it does not satisfy this assert
+    # assert (output.size(1) == (5 + num_classes) * num_anchors)
+
+    # print(output.size())
+
+    # Slice the second dimension (channel) of output into:
+    # [ 2, 2, 1, num_classes, 2, 2, 1, num_classes, 2, 2, 1, num_classes ]
+    # And then into
+    # bxy = [ 6 ] bwh = [ 6 ] det_conf = [ 3 ] cls_conf = [ num_classes * 3 ]
+    # batch = output.size(0)
+    # H = output.size(2)
+    # W = output.size(3)
+
+    bxy_list = []
+    bwh_list = []
+    det_confs_list = []
+    cls_confs_list = []
+
+    for i in range(num_anchors):
+        begin = i * (5 + num_classes)
+        end = (i + 1) * (5 + num_classes)
+        
+        bxy_list.append(output[:, begin : begin + 2])
+        bwh_list.append(output[:, begin + 2 : begin + 4])
+        det_confs_list.append(output[:, begin + 4 : begin + 5])
+        cls_confs_list.append(output[:, begin + 5 : end])
+
+    # Shape: [batch, num_anchors * 2, H, W]
+    bxy = torch.cat(bxy_list, dim=1)
+    # Shape: [batch, num_anchors * 2, H, W]
+    bwh = torch.cat(bwh_list, dim=1)
+
+    # Shape: [batch, num_anchors, H, W]
+    det_confs = torch.cat(det_confs_list, dim=1)
+    # Shape: [batch, num_anchors * H * W]
+    det_confs = det_confs.view(output.size(0), num_anchors * output.size(2) * output.size(3))
+
+    # Shape: [batch, num_anchors * num_classes, H, W]
+    cls_confs = torch.cat(cls_confs_list, dim=1)
+    # Shape: [batch, num_anchors, num_classes, H * W]
+    cls_confs = cls_confs.view(output.size(0), num_anchors, num_classes, output.size(2) * output.size(3))
+    # Shape: [batch, num_anchors, num_classes, H * W] --> [batch, num_anchors * H * W, num_classes] 
+    cls_confs = cls_confs.permute(0, 1, 3, 2).reshape(output.size(0), num_anchors * output.size(2) * output.size(3), num_classes)
+
+    # Apply sigmoid(), exp() and softmax() to slices
+    #
+    bxy = torch.sigmoid(bxy) * scale_x_y - 0.5 * (scale_x_y - 1)
+    bwh = torch.exp(bwh)
+    det_confs = torch.sigmoid(det_confs)
+    cls_confs = torch.sigmoid(cls_confs)
+
+    # Prepare C-x, C-y, P-w, P-h (None of them are torch related)
+    grid_x = np.expand_dims(np.expand_dims(np.expand_dims(np.linspace(0, output.size(3) - 1, output.size(3)), axis=0).repeat(output.size(2), 0), axis=0), axis=0)
+    grid_y = np.expand_dims(np.expand_dims(np.expand_dims(np.linspace(0, output.size(2) - 1, output.size(2)), axis=1).repeat(output.size(3), 1), axis=0), axis=0)
+    # grid_x = torch.linspace(0, W - 1, W).reshape(1, 1, 1, W).repeat(1, 1, H, 1)
+    # grid_y = torch.linspace(0, H - 1, H).reshape(1, 1, H, 1).repeat(1, 1, 1, W)
+
+    anchor_w = []
+    anchor_h = []
+    for i in range(num_anchors):
+        anchor_w.append(anchors[i * 2])
+        anchor_h.append(anchors[i * 2 + 1])
+
+    device = None
+    cuda_check = output.is_cuda
+    if cuda_check:
+        device = output.get_device()
+
+    bx_list = []
+    by_list = []
+    bw_list = []
+    bh_list = []
+
+    # Apply C-x, C-y, P-w, P-h
+    for i in range(num_anchors):
+        ii = i * 2
+        # Shape: [batch, 1, H, W]
+        bx = bxy[:, ii : ii + 1] + torch.tensor(grid_x, device=device, dtype=torch.float32) # grid_x.to(device=device, dtype=torch.float32)
+        # Shape: [batch, 1, H, W]
+        by = bxy[:, ii + 1 : ii + 2] + torch.tensor(grid_y, device=device, dtype=torch.float32) # grid_y.to(device=device, dtype=torch.float32)
+        # Shape: [batch, 1, H, W]
+        bw = bwh[:, ii : ii + 1] * anchor_w[i]
+        # Shape: [batch, 1, H, W]
+        bh = bwh[:, ii + 1 : ii + 2] * anchor_h[i]
+
+        bx_list.append(bx)
+        by_list.append(by)
+        bw_list.append(bw)
+        bh_list.append(bh)
+
+
+    ########################################
+    #   Figure out bboxes from slices     #
+    ########################################
+    
+    # Shape: [batch, num_anchors, H, W]
+    bx = torch.cat(bx_list, dim=1)
+    # Shape: [batch, num_anchors, H, W]
+    by = torch.cat(by_list, dim=1)
+    # Shape: [batch, num_anchors, H, W]
+    bw = torch.cat(bw_list, dim=1)
+    # Shape: [batch, num_anchors, H, W]
+    bh = torch.cat(bh_list, dim=1)
+
+    # Shape: [batch, 2 * num_anchors, H, W]
+    bx_bw = torch.cat((bx, bw), dim=1)
+    # Shape: [batch, 2 * num_anchors, H, W]
+    by_bh = torch.cat((by, bh), dim=1)
+
+    # normalize coordinates to [0, 1]
+    bx_bw /= output.size(3)
+    by_bh /= output.size(2)
+
+    # Shape: [batch, num_anchors * H * W, 1]
+    bx = bx_bw[:, :num_anchors].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
+    by = by_bh[:, :num_anchors].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
+    bw = bx_bw[:, num_anchors:].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
+    bh = by_bh[:, num_anchors:].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
+
+    bx1 = bx - bw * 0.5
+    by1 = by - bh * 0.5
+    bx2 = bx1 + bw
+    by2 = by1 + bh
+
+    # Shape: [batch, num_anchors * h * w, 4] -> [batch, num_anchors * h * w, 1, 4]
+    boxes = torch.cat((bx1, by1, bx2, by2), dim=2).view(output.size(0), num_anchors * output.size(2) * output.size(3), 1, 4)
+    # boxes = boxes.repeat(1, 1, num_classes, 1)
+
+    # boxes:     [batch, num_anchors * H * W, 1, 4]
+    # cls_confs: [batch, num_anchors * H * W, num_classes]
+    # det_confs: [batch, num_anchors * H * W]
+
+    det_confs = det_confs.view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
+    confs = cls_confs * det_confs
+
+    # boxes: [batch, num_anchors * H * W, 1, 4]
+    # confs: [batch, num_anchors * H * W, num_classes]
+
+    return  boxes, confs
+
+class YoloLayer(nn.Module):
+    ''' Yolo layer
+    model_out: while inference,is post-processing inside or outside the model
+        true:outside
+    '''
+    def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, stride=32, model_out=False):
+        super(YoloLayer, self).__init__()
+        self.anchor_mask = anchor_mask
+        self.num_classes = num_classes
+        self.anchors = anchors
+        self.num_anchors = num_anchors
+        self.anchor_step = len(anchors) // num_anchors
+        self.coord_scale = 1
+        self.noobject_scale = 1
+        self.object_scale = 5
+        self.class_scale = 1
+        self.thresh = 0.6
+        self.stride = stride
+        self.seen = 0
+        self.scale_x_y = 1
+
+        self.model_out = model_out
+
+    def forward(self, output, target=None):
+        if self.training:
+            return output
+        masked_anchors = []
+        for m in self.anchor_mask:
+            masked_anchors += self.anchors[m * self.anchor_step:(m + 1) * self.anchor_step]
+        masked_anchors = [anchor / self.stride for anchor in masked_anchors]
+
+        return yolo_forward_dynamic(output, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask),scale_x_y=self.scale_x_y)
+

install.bat

@@ -12,16 +12,26 @@ BATCH
 echo '**** GIMP-ML Setup started ****'
 if python --version 2>&1 | grep -q '^Python 3\.'; then #
     echo 'Python 3 found.' #
+    sudo apt-get install python3-distutils #
+    sudo apt install python3-testresources #
+    wget https://bootstrap.pypa.io/pip/get-pip.py #
+    python GIMP-ML/get-pip.py #
+    rm -rf get-pip* #
     python -m pip install virtualenv
     python -m virtualenv gimpenv3 #
     source gimpenv3/bin/activate #
-    python -m pip3 install torch torchvision -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html #
+    python -m pip install torch torchvision -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html #
     python -m pip install GIMP-ML/.
     python -c "import gimpml; gimpml.setup_python_weights()" #
     chmod -R 777 gimpenv3/lib/ #
     deactivate #
 elif python3 --version 2>&1 | grep -q '^Python 3\.'; then #
     echo 'Python 3 found.' #
+    sudo apt-get install python3-distutils #
+    sudo apt install python3-testresources #
+    wget https://bootstrap.pypa.io/pip/get-pip.py #
+    python3 get-pip.py #
+    rm -rf get-pip* #
     python3 -m pip install virtualenv
     python3 -m virtualenv gimpenv3 #
     source gimpenv3/bin/activate #

testscases/test.py

@@ -4,6 +4,9 @@ import gimpml
 image = cv2.imread("sampleinput/img.png")
 alpha = cv2.imread("sampleinput/alpha.png")
 
+out = gimpml.edge(image)
+cv2.imwrite("output/tmp-edge.jpg", out)
+
 # out = gimpml.kmeans(image)
 # cv2.imwrite("output/tmp-kmeans.jpg", out)
 # #