Refactor

celldetection/models/__init__.py

@@ -1,11 +1,7 @@
-from .commons import TwoConvBnRelu
-from .unet import UNetEncoder, UNet, U12, U17, U22, SlimU22, WideU22
-from .resnet import get_resnet, ResNet50, ResNet34, ResNet18, ResNet152, ResNet101, WideResNet101_2, WideResNet50_2, \
-    ResNeXt152_32x8d, ResNeXt101_32x8d, ResNeXt50_32x4d
-from .cpn import CPN, CpnSlimU22, CpnU22, CpnWideU22, CpnResNet18FPN, CpnResNet34FPN, CpnResNet50FPN, CpnResNet101FPN, \
-    CpnResNet152FPN, CpnResNeXt50FPN, CpnResNeXt101FPN, CpnResNeXt152FPN, CpnWideResNet50FPN, \
-    CpnWideResNet101FPN, CpnMobileNetV3LargeFPN, CpnMobileNetV3SmallFPN
-from .fpn import FPN, ResNeXt50FPN, ResNeXt101FPN, ResNet18FPN, ResNet34FPN, ResNeXt152FPN, WideResNet50FPN, \
-    WideResNet101FPN, ResNet50FPN, ResNet101FPN, ResNet152FPN, MobileNetV3SmallFPN, MobileNetV3LargeFPN
-from .inference import Inference
+from .commons import *
+from .unet import *
+from .resnet import *
+from .cpn import *
+from .fpn import *
+from .inference import *
 from .mobilenetv3 import *

celldetection/models/commons.py

@@ -1,5 +1,7 @@
 import torch.nn as nn
 
+__all__ = ['TwoConvBnRelu']
+
 
 class TwoConvBnRelu(nn.Sequential):
     def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, stride=1, mid_channels=None, **kwargs):

celldetection/models/cpn.py

@@ -13,6 +13,10 @@
 from .fpn import ResNet34FPN, ResNet18FPN, ResNet50FPN, ResNet101FPN, ResNet152FPN, ResNeXt50FPN, \
     ResNeXt101FPN, ResNeXt152FPN, WideResNet50FPN, WideResNet101FPN, MobileNetV3LargeFPN, MobileNetV3SmallFPN
 
+__all__ = ['CPN', 'CpnSlimU22', 'CpnU22', 'CpnWideU22', 'CpnResNet18FPN', 'CpnResNet34FPN', 'CpnResNet50FPN',
+           'CpnResNet101FPN', 'CpnResNet152FPN', 'CpnResNeXt50FPN', 'CpnResNeXt101FPN', 'CpnResNeXt152FPN',
+           'CpnWideResNet50FPN', 'CpnWideResNet101FPN', 'CpnMobileNetV3LargeFPN', 'CpnMobileNetV3SmallFPN']
+
 
 class ReadOut(nn.Module):
     def __init__(

celldetection/models/fpn.py

@@ -4,6 +4,10 @@
     ResNeXt152_32x8d, WideResNet50_2, WideResNet101_2
 from .mobilenetv3 import MobileNetV3Large, MobileNetV3Small
 
+__all__ = ['FPN', 'ResNeXt50FPN', 'ResNeXt101FPN', 'ResNet18FPN', 'ResNet34FPN', 'ResNeXt152FPN', 'WideResNet50FPN',
+           'WideResNet101FPN', 'ResNet50FPN', 'ResNet101FPN', 'ResNet152FPN', 'MobileNetV3SmallFPN',
+           'MobileNetV3LargeFPN']
+
 
 class FPN(BackboneWithFPN):
     def __init__(self, backbone, channels=256, return_layers: dict = None):

celldetection/models/inference.py

@@ -1,6 +1,8 @@
 import torch
 from ..util.util import asnumpy
 
+__all__ = ['Inference']
+
 
 class Inference:
     def __init__(self, model, device=None, amp=False, transforms=None):

celldetection/models/resnet.py

@@ -4,6 +4,9 @@
 from torchvision.models.resnet import ResNet as RN, Bottleneck, BasicBlock
 from ..util.util import Dict
 
+__all__ = ['get_resnet', 'ResNet50', 'ResNet34', 'ResNet18', 'ResNet152', 'ResNet101', 'WideResNet101_2',
+           'WideResNet50_2', 'ResNeXt152_32x8d', 'ResNeXt101_32x8d', 'ResNeXt50_32x4d']
+
 
 def make_res_layer(block, inplanes, planes, blocks, norm_layer=nn.BatchNorm2d, base_width=64, groups=1, stride=1,
                    dilation=1, dilate=False, **kwargs) -> nn.Module:

celldetection/models/unet.py

@@ -6,9 +6,10 @@
 from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork
 from collections import OrderedDict
 from typing import List, Tuple, Dict
-
 from .commons import TwoConvBnRelu
 
+__all__ = ['UNetEncoder', 'UNet', 'U12', 'U17', 'U22', 'SlimU22', 'WideU22']
+
 
 class UNetEncoder(nn.Sequential):
     def __init__(self, in_channels, depth=5, base_channels=64, factor=2, pool=True):

celldetection/ops/__init__.py

@@ -0,0 +1 @@
+from .commons import *

celldetection/ops/commons.py

@@ -2,6 +2,8 @@
 import torch.nn.functional as F
 from typing import List
 
+__all__ = ['downsample_labels']
+
 
 def downsample_labels(inputs, size: List[int]):
     """

celldetection/util/__init__.py

@@ -1,3 +1,2 @@
-from .util import Dict, lookup_nn, reduce_loss_dict, to_device, asnumpy, fetch_model, random_code_name, dict_hash, \
-    fetch_image, random_seed
-from .timer import start_timer, stop_timer, print_timing
+from .util import *
+from .timer import *

celldetection/util/timer.py

@@ -3,6 +3,8 @@
 from time import time
 import numpy as np
 
+__all__ = ['start_timer', 'stop_timer', 'print_timing']
+
 TIMINGS = {}
 
 
@@ -26,7 +28,8 @@ def convert_seconds(seconds):
 
 
 def seconds_to_str(seconds):
-    s = [f'{i} {n[:-1] if i == 1 else n}' for i, n in zip(convert_seconds(seconds), ('days', 'hours', 'minutes', 'seconds'))]
+    s = [f'{i} {n[:-1] if i == 1 else n}' for i, n in
+         zip(convert_seconds(seconds), ('days', 'hours', 'minutes', 'seconds'))]
     s = ', '.join(s)
     return s
 

celldetection/util/util.py

@@ -8,6 +8,9 @@
 import hashlib
 import json
 
+__all__ = ['Dict', 'lookup_nn', 'reduce_loss_dict', 'to_device', 'asnumpy', 'fetch_model', 'random_code_name',
+           'dict_hash', 'fetch_image', 'random_seed']
+
 
 class Dict(dict):
     __getattr__ = dict.__getitem__  # alternative: dict.get if KeyError is not desired