rm models weights, fix DL interfaces (#74)

README.rst

@@ -62,6 +62,7 @@ The repository includes the following directories:
 * Package `cases <https://github.com/ITMO-NSS-team/GEFEST/tree/main/cases>`__ includes several *how-to-use-cases* where you can start to discover how GEFEST works;
 * All *unit and integration tests* can be observed in the `test <https://github.com/ITMO-NSS-team/GEFEST/tree/main/test>`__ directory;
 * The sources of the documentation are in the `docs <https://github.com/ITMO-NSS-team/GEFEST/tree/main/docs>`__.
+* Weights of pretrained DL models can be downloaded from `this repository <https://gitlab.actcognitive.org/itmo-nss-team/gefest-models>`__.
 
 Cases and examples
 ==================

gefest/tools/estimators/DL/bw_surrogate/bw_cnn.py

@@ -1,28 +1,22 @@
 import os
 import shutil
-from typing import Optional
 
 import matplotlib
 import matplotlib.pyplot as plt
 import tensorflow as tf
 from tensorflow import keras
 
-from gefest.core.structure.domain import Domain
-from gefest.core.structure.structure import Structure
+from gefest.core.geometry import Structure
+from gefest.core.geometry.domain import Domain
+from gefest.tools.estimators.estimator import Estimator
 
 matplotlib.use('agg')
 
 
-class BWCNN:
-    """
-    ::TODO:: Make abstract version to create own realizations for specific tasks
-    """
+class BWCNN(Estimator):
+    """Surrogate model for breakwaters task."""
 
-    """
-    Surrogate model for breakwaters task
-    """
-
-    def __init__(self, path, domain: Domain, main_model: Optional = None):
+    def __init__(self, path, domain: Domain, main_model=None):
         super(BWCNN, self).__init__()
 
         self.domain = domain
@@ -35,24 +29,22 @@ def __init__(self, path, domain: Domain, main_model: Optional = None):
         self.rate = 4
 
     def _create_temp_path(self):
-        """
-        Creation of temporary folder for images
-        :return: None
-        """
+        """Creation of temporary folder for images."""
         path = 'tmp_images'
 
         if os.path.exists(path):
             shutil.rmtree(path)
+
         os.makedirs(path)
 
         return
 
     def _save_as_fig(self, struct: Structure, ax=plt):
-        """
-        Saving structure as image
-        :param struct: (Structure)
-        :param ax: figure
-        :return: None
+        """Saves structs as images.
+
+        Args:
+            struct (Structure): _description_
+            ax : Plot. Defaults to plt.
         """
         plt.style.use('dark_background')
 
@@ -82,10 +74,13 @@ def _save_as_fig(self, struct: Structure, ax=plt):
         ax.close('all')
 
     def _to_tensor(self, struct: Structure):
-        """
-        Transformation structure to binary tensor
-        :param struct: (Structure), input structure
-        :return: (Tensor), binary matrix with WxHx1 dimension
+        """Transformation structure to binary tensor.
+
+        Args:
+            struct (Structure): Input structure
+
+        Returns:
+            Tensor: Binary matrix with WxHx1 dimension.
         """
         self._save_as_fig(struct)
 
@@ -97,10 +92,13 @@ def _to_tensor(self, struct: Structure):
         return image_tensor
 
     def estimate(self, struct: Structure):
-        """
-        Estimation step
-        :param struct: (Structure), input structure
-        :return: (Float), performance of structure
+        """Estimation step.
+
+        Args:
+            struct (Structure), input structure.
+
+        Returns:
+            (float): Performance.
         """
         tensor = self._to_tensor(struct)
         tensor = tf.reshape(tensor, (1, self.img_size, self.img_size, 1))

gefest/tools/estimators/DL/heat/heat_cnn.py

@@ -7,15 +7,11 @@
 from torch import nn
 from torchvision import transforms
 
+from gefest.tools.estimators.estimator import Estimator
 
-class HeatCNN:
-    """
-    ::TODO:: Make abstract version for creation own realizations for specific tasks
-    """
 
-    """
-    Surrogate model for the heat components task
-    """
+class HeatCNN(Estimator):
+    """Surrogate model for the heat components task."""
 
     def __init__(self, path):
         super(HeatCNN, self).__init__()
@@ -26,10 +22,13 @@ def __init__(self, path):
         self.img_size = 128
 
     def estimate(self, obj):
-        """
-        Estimation step
-        :param obj: (Tensor), [1 x C x W x H], object for estimate
-        :return: (Float), performance of object
+        """Estimation step.
+
+        Args:
+            obj (torch.Tensor): [1 x C x W x H], object for estimate.
+
+        Returns:
+            (Float): performance of object.
         """
         plt.imsave('tmp_images/0.png', torch.squeeze(torch.tensor(obj)))
 
@@ -50,6 +49,7 @@ def estimate(self, obj):
 
 
 class EffModel(nn.Module):
+    """Efficient net surrogate model."""
     def __init__(self):
         super(EffModel, self).__init__()
 

gefest/tools/samplers/DL/heat/heat_sampler.py

@@ -1,6 +1,7 @@
 import os
 
 import torch
+from tools.samplers.sampler import Sampler
 
 from gefest.tools.samplers.DL.microfluid.aae import AAE
 from gefest.tools.samplers.DL.microfluid.backbones import (
@@ -12,10 +13,10 @@
 os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 
 
-class DeepSampler:
-    """
-    Deep learning sampler for microfluidic problem based on adversarial auto encoder.
-    It is creates images of polygons with size 128x128
+class DeepSampler(Sampler):
+    """Deep learning sampler for microfluidic problem based on adversarial auto encoder.
+
+    It is creates images of polygons with size 128x128.
     """
 
     def __init__(self, path):
@@ -30,10 +31,7 @@ def __init__(self, path):
         self._configurate_sampler()
 
     def _configurate_sampler(self):
-        """
-        Configurate deep sampler using configuration parameters
-        :return: None
-        """
+        """Configurate deep sampler using configuration parameters."""
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         conv_dims = [32, 64, 128, 256, 512, 512]  # We define 6 layers encoder and decoder
         n_layers = 2
@@ -57,13 +55,16 @@ def _configurate_sampler(self):
         self.sampler = aae
 
     def sample(self, n_samples: int, domain):
-        """
-        Sampling procedure using deep learning sampler.
+        """Sampling procedure using deep learning sampler.
+
         It based on general GEFEST deep learning sampler architecture,
-        i.e. on mapping noise to object
+        i.e. on mapping noise to object.
+
+        Args:
+            n_samples (Int): Number of samples.
 
-        :param n_samples: (Int) number of samples
-        :return: (List(Structure)) sample n_samples structures
+        Returns:
+            (List(Structure)): Sample n_samples structures.
         """
         with torch.no_grad():
             noise = torch.normal(mean=0, std=1, size=(n_samples, self.hidden_dim)).to(self.device)

gefest/tools/samplers/DL/microfluid/microfluid_sampler.py

@@ -3,6 +3,7 @@
 import cv2 as cv
 import numpy as np
 import torch
+from tools.samplers.sampler import Sampler
 
 from gefest.core.geometry import Point, Polygon, Structure
 from gefest.core.geometry.domain import Domain
@@ -16,16 +17,13 @@
 os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 
 
-class DeepSampler:
-    """
-    Deep learning sampler for microfluidic problem based on adversarial auto encoder.
-    It is creates images of polygons with size 128x128
+class DeepSampler(Sampler):
+    """Deep learning sampler for microfluidic problem based on adversarial auto encoder.
+
+    It is creates images of polygons with size 128x128.
     """
 
     def __init__(self, path):
-        """
-        :param path: path to deep learning generative model
-        """
         super(DeepSampler, self).__init__()
 
         self.path = path
@@ -37,12 +35,9 @@ def __init__(self, path):
         self._configurate_sampler()
 
     def _configurate_sampler(self):
-        """
-        Configurate deep sampler using configuration parameters
-        :return: None
-        """
+        """Configurate deep sampler using configuration parameters."""
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        conv_dims = [32, 64, 128, 256, 256, 512]  # We define 6 layers encoder and decoder
+        conv_dims = [32, 64, 128, 256, 256, 512]  # 6 layers
         n_layers = 2
         self.hidden_dim = 32
 
@@ -66,15 +61,12 @@ def _configurate_sampler(self):
     def _transform(self, objects, domain) -> list[Structure]:
         """Transforms images to polygons using edge detector.
 
-        :param objects: (Array) [n_samples x 1 x 128 x 128]
-        :return: List(Structure)
-
         Args:
-            objects (_type_): _description_
-            domain (_type_): _description_
+            objects (Array): [n_samples x 1 x 128 x 128].
+            domain (Doamin): Task domain.
 
         Returns:
-            _type_: _description_
+            List(Structure): Structures.
         """
         samples = []
 
@@ -110,15 +102,7 @@ def _transform(self, objects, domain) -> list[Structure]:
         return samples
 
     def sample(self, n_samples: int, domain: Domain):
-        """
-        Sampling procedure using deep learning sampler.
-        It based on general GEFEST deep learning sampler architecture,
-        i.e. on mapping noise to object
-
-        :param n_samples: (Int) number of samples
-        :param domain: (Domain) design domain
-        :return: (List(Structure)) sample n_samples structures
-        """
+        """Sampling procedure using deep learning sampler."""
         with torch.no_grad():
             noise = torch.normal(mean=0, std=1, size=(n_samples, self.hidden_dim)).to(self.device)
             objects = self.sampler.decoder.sample(noise).numpy()  # Numpy: {n_samples, 1, 128, 128}