easymode 0.0.1__tar.gz

easymode-0.0.1/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: easymode
+Version: 0.0.1
+Summary: Easymode - a collection of pretrained general networks for segmenting common eukaryotic features in cryoET
+Home-page: 
+Author: mgflast
+Author-email: mgflast@gmail.com
+License: GPL v3
+Requires-Dist: mrcfile
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: huggingface_hub
+Requires-Dist: requests
+Requires-Dist: tifffile
+Requires-Dist: psutil
+Dynamic: author
+Dynamic: author-email
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: summary

easymode-0.0.1/easymode/core/augmentations.py

@@ -0,0 +1,53 @@
+import numpy as np
+import random
+from scipy.ndimage import rotate, gaussian_filter, median_filter
+
+from easymode.membrain_fourier_augmentations.fourier_augmentations import MissingWedgeMaskAndFourierAmplitudeMatchingCombined
+
+ROT_XZ_YZ_MAX_ANGLE = 10.0
+ROT_XY_MAX_ANGLE = 10.0
+
+def rotate_90_xy(img, label):
+    k = random.randint(0, 3)
+    img = np.rot90(img, k=k, axes=(1, 2))
+    label = np.rot90(label, k=k, axes=(1, 2))
+    return img, label
+
+def rotate_90_xz(img, label):
+    k = random.randint(0, 1) * 2
+    img = np.rot90(img, k=k, axes=(0, 2))
+    label = np.rot90(label, k=k, axes=(0, 2))
+    return img, label
+
+def flip(img, label):
+    k = random.choice([None, 0, 1, 2])
+    if k is not None:
+        img = np.flip(img, axis=k)
+        label = np.flip(label, axis=k)
+    return img, label
+
+def rotate_continuous_xz_or_yz(img, label):
+    plane = random.choice([(0, 2), (0, 1)])
+    angle = np.random.uniform(-ROT_XZ_YZ_MAX_ANGLE, ROT_XZ_YZ_MAX_ANGLE)
+
+    img = rotate(img, angle, axes=plane, order=1, mode='reflect', prefilter=False, reshape=False)
+    label = rotate(label, angle, axes=plane, order=0, mode='constant', cval=2, reshape=False)
+
+    return img, label
+
+def rotate_continuous_xy(img, label):
+    angle = np.random.uniform(-ROT_XY_MAX_ANGLE, ROT_XY_MAX_ANGLE)
+
+    img = rotate(img, angle, axes=(1, 2), order=1, mode='reflect', prefilter=False, reshape=False)
+    label = rotate(label, angle, axes=(1, 2), order=0, mode='constant', cval=2, reshape=False)
+
+    return img, label
+
+def remove_wedge(img, label):
+    membrain_fourier_trickery_machine = MissingWedgeMaskAndFourierAmplitudeMatchingCombined()
+    img = membrain_fourier_trickery_machine(img)
+    return img, label
+
+def filter_gaussian(img, label):
+    img = gaussian_filter(img, sigma=random.uniform(0.3, 1.5))
+    return img, label

easymode-0.0.1/easymode/core/config.py

@@ -0,0 +1,49 @@
+import os, shutil, json
+
+version = "0.0.1"
+license = "GNU GPL v3"
+
+root = os.path.dirname(os.path.dirname(__file__))
+settings_path = os.path.join(os.path.expanduser("~"), "easymode", "settings.txt")
+
+
+def parse_settings():
+    # If settings file not found, copy the one from core to the right location.
+    if not os.path.exists(settings_path):
+        os.makedirs(os.path.dirname(settings_path), exist_ok=True)
+        shutil.copy(os.path.join(root, "core", "settings.txt"), settings_path)
+
+    try:
+        with open(settings_path, 'r') as f:
+            sdict = json.load(f)
+    except Exception as e:
+        shutil.copy(os.path.join(root, "core", "settings.txt"), settings_path)
+        parse_settings()
+        return
+
+    # Read settings - if any parameters are missing, insert them.
+    with open(os.path.join(root, "core", "settings.txt"), 'r') as f:
+        default_settings = json.load(f)
+
+    for key in default_settings:
+        if key not in sdict:
+            sdict[key] = default_settings[key]
+
+    if sdict["MODEL_DIRECTORY"] == "" or not os.path.exists(sdict["MODEL_DIRECTORY"]):
+        sdict["MODEL_DIRECTORY"] = os.path.join(os.path.expanduser("~"), "easymode")
+
+    with open(settings_path, 'w') as f:
+        json.dump(sdict, f, indent=2)
+
+    return sdict
+
+
+settings = parse_settings()
+
+
+def edit_setting(key, value):
+    global settings
+    settings[key] = value
+    with open(settings_path, 'w') as f:
+        json.dump(settings, f, indent=2)
+    print(key, value)

easymode-0.0.1/easymode/core/distribution.py

@@ -0,0 +1,206 @@
+import os
+import requests
+import json
+from pathlib import Path
+from huggingface_hub import hf_hub_download, HfApi
+from easymode.core.model import create
+import tensorflow as tf
+import easymode.core.config as cfg
+
+# Configuration
+HF_REPO_ID = "mgflast/easymode"  # Single repo for all models
+MODEL_CACHE_DIR = cfg.settings["MODEL_DIRECTORY"]
+VERSION_FILE = "model_info.json"
+
+
+def get_model_info(model_title):
+    """Get model repository and filename info."""
+    filename = f"{model_title}_3d.h5"
+
+    return {
+        'repo_id': HF_REPO_ID,
+        'filename': filename,
+        'local_path': os.path.join(MODEL_CACHE_DIR, filename),
+        'version_path': os.path.join(MODEL_CACHE_DIR, f"{model_title}_3d_info.json")
+    }
+
+
+def is_online():
+    """Check if internet connection is available."""
+    try:
+        response = requests.get("https://huggingface.co", timeout=5)
+        return response.status_code == 200
+    except:
+        return False
+
+
+def get_remote_version(repo_id):
+    """Get latest version info from Hugging Face."""
+    try:
+        api = HfApi()
+        repo_info = api.repo_info(repo_id)
+
+        # Use last modified time as version identifier
+        last_modified = repo_info.last_modified
+        return {
+            'version': last_modified.isoformat() if last_modified else "unknown",
+            'commit_hash': repo_info.sha[:8] if repo_info.sha else "unknown"
+        }
+    except Exception as e:
+        print(f"Warning: Could not get remote version info: {e}")
+        return None
+
+
+def get_local_version(version_path):
+    """Get local version info."""
+    if not os.path.exists(version_path):
+        return None
+
+    try:
+        with open(version_path, 'r') as f:
+            return json.load(f)
+    except:
+        return None
+
+
+def save_version_info(version_path, version_info):
+    """Save version info locally."""
+    os.makedirs(os.path.dirname(version_path), exist_ok=True)
+    with open(version_path, 'w') as f:
+        json.dump(version_info, f, indent=2)
+
+
+def download_model(repo_id, filename, local_path, version_path):
+    """Download model from Hugging Face."""
+    print(f"Downloading {repo_id}/{filename}...")
+
+    try:
+        # Ensure cache directory exists
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+
+        # Download file
+        downloaded_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            cache_dir=MODEL_CACHE_DIR,
+            local_dir=os.path.dirname(local_path)
+        )
+
+        # Get and save version info
+        remote_version = get_remote_version(repo_id)
+        if remote_version:
+            save_version_info(version_path, remote_version)
+
+        print(f"\nDownloaded successfully to {local_path}")
+        return local_path
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to download {repo_id}: {e}")
+
+
+def load_model_weights(weights_path):
+    model = create()
+    dummy_input = tf.zeros((1, 160, 160, 160, 1))
+    _ = model(dummy_input)
+    model.load_weights(weights_path)
+    return model
+
+def cache_model(model_title, force_download=False, silent=False):
+    info = get_model_info(model_title)
+    online = is_online()
+
+    # Check if local file exists
+    local_exists = os.path.exists(info['local_path'])
+
+    if force_download or not local_exists:
+        # Need to download
+        if not online:
+            if local_exists:
+                print("Local model found. There may be updates available, but we cannot check without an internet connection.")
+            else:
+                print(f"The required network weights are not available in the local cache {MODEL_CACHE_DIR} and there is no internet connection available to download them - aborting...")
+                exit()
+        else:
+            print(f"The required network weights for {model_title} are not available in the local cache. Downloading now...")
+            download_model(info['repo_id'], info['filename'],
+                           info['local_path'], info['version_path'])
+
+    elif local_exists and online:
+        # Check if we need to update
+        local_version = get_local_version(info['version_path'])
+        remote_version = get_remote_version(info['repo_id'])
+
+        if remote_version and local_version:
+            if remote_version['version'] != local_version['version']:
+                if not silent:
+                    print(f"New version available for {model_title}, updating...")
+                download_model(info['repo_id'], info['filename'], info['local_path'], info['version_path'])
+        elif remote_version and not local_version:
+            if not silent:
+                print("No local version info, checking for updates...")
+            save_version_info(info['version_path'], remote_version)
+
+    return info['local_path']
+
+def load_model(local_path):
+    return load_model_weights(local_path)
+
+def clear_model_cache(model_title=None):
+    """Clear local model cache."""
+    if model_title:
+        # Clear specific model
+        info = get_model_info(model_title)
+        files_to_remove = [info['local_path'], info['version_path']]
+
+        for file_path in files_to_remove:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Removed {file_path}")
+    else:
+        # Clear all models
+        import shutil
+        if os.path.exists(MODEL_CACHE_DIR):
+            shutil.rmtree(MODEL_CACHE_DIR)
+            print(f"Cleared model cache: {MODEL_CACHE_DIR}")
+
+def list_remote_models():
+    """List all available models in the Hugging Face repository."""
+    if not is_online():
+        print("Cannot list remote models: No internet connection")
+        return []
+
+    try:
+        api = HfApi()
+        repo_files = api.list_repo_files(HF_REPO_ID)
+
+        # Filter for .h5 model files
+        model_files = [f for f in repo_files if f.endswith('.h5')]
+
+        if not model_files:
+            print("No model files found in repository")
+            return []
+
+        print()
+        print(f"Easymode can currently segment the following features:")
+        print()
+        models = []
+
+        for model_file in sorted(model_files):
+            model_name = model_file.replace('.h5', '')
+            if '_' in model_name:
+                title, dim = model_name.rsplit('_', 1)
+                models.append({'title': title, 'filename': model_file})
+
+                # Check if we have it locally
+                local_path = os.path.join(MODEL_CACHE_DIR, model_file)
+                local_status = "weights local" if os.path.exists(local_path) else "weights available for download"
+
+                print(f"   > {title} - [{local_status}]")
+            else:
+                pass
+
+        return models
+
+    except Exception as e:
+        print(f"Error listing remote models: {e}")
+        return []

easymode-0.0.1/easymode/core/inference.py

@@ -0,0 +1,257 @@
+import os, glob, time, multiprocessing, psutil
+import tensorflow as tf
+import gc
+from tensorflow.keras import mixed_precision
+import mrcfile
+import numpy as np
+from easymode.core.distribution import cache_model, load_model
+
+TILE_SIZE = 160
+OVERLAP = 32
+MAX_CHUNK_SIZE = 64
+
+def tile_volume(volume, patch_size=TILE_SIZE, overlap=OVERLAP):
+    d, h, w = volume.shape
+    stride = patch_size - 2 * overlap
+
+    z_boxes = max(1, (d + stride - 1) // stride)
+    y_boxes = max(1, (h + stride - 1) // stride)
+    x_boxes = max(1, (w + stride - 1) // stride)
+
+    tiles = []
+    positions = []
+
+    for z_idx in range(z_boxes):
+        for y_idx in range(y_boxes):
+            for x_idx in range(x_boxes):
+                z_start = z_idx * stride - overlap
+                y_start = y_idx * stride - overlap
+                x_start = x_idx * stride - overlap
+
+                vol_z_start = max(0, z_start)
+                vol_y_start = max(0, y_start)
+                vol_x_start = max(0, x_start)
+
+                vol_z_end = min(d, z_start + patch_size)
+                vol_y_end = min(h, y_start + patch_size)
+                vol_x_end = min(w, x_start + patch_size)
+
+                extracted = volume[vol_z_start:vol_z_end, vol_y_start:vol_y_end, vol_x_start:vol_x_end]
+
+                tile = np.zeros((patch_size, patch_size, patch_size), dtype=volume.dtype)
+
+                tile_z_start = vol_z_start - z_start
+                tile_y_start = vol_y_start - y_start
+                tile_x_start = vol_x_start - x_start
+
+                tile[tile_z_start:tile_z_start + extracted.shape[0],
+                tile_y_start:tile_y_start + extracted.shape[1],
+                tile_x_start:tile_x_start + extracted.shape[2]] = extracted
+
+                tiles.append(tile)
+                positions.append((z_idx * stride, y_idx * stride, x_idx * stride))
+
+    tiles = np.array(tiles)
+    tiles = np.expand_dims(tiles, axis=-1)
+
+    return tiles, positions, volume.shape
+
+
+def detile_volume(segmented_tiles, positions, original_shape, patch_size=TILE_SIZE, overlap=OVERLAP):
+    d, h, w = original_shape
+    output_volume = np.zeros((d, h, w), dtype=np.float32)
+    stride = patch_size - 2 * overlap
+
+    if segmented_tiles.ndim == 5:
+        segmented_tiles = segmented_tiles.squeeze(-1)
+
+    for tile, (z_pos, y_pos, x_pos) in zip(segmented_tiles, positions):
+        center_region = tile[overlap:overlap + stride, overlap:overlap + stride, overlap:overlap + stride]
+
+        z_end = min(z_pos + stride, d)
+        y_end = min(y_pos + stride, h)
+        x_end = min(x_pos + stride, w)
+
+        actual_z = z_end - z_pos
+        actual_y = y_end - y_pos
+        actual_x = x_end - x_pos
+
+        output_volume[z_pos:z_end, y_pos:y_end, x_pos:x_end] = center_region[:actual_z, :actual_y, :actual_x]
+
+    return output_volume
+
+
+def create_weight_matrix(patch_size, border_size=OVERLAP):
+    """Create weight matrix that reduces border artifacts."""
+    weights = np.ones((patch_size, patch_size, patch_size), dtype=np.float32)
+    fade_size = 16
+
+    for dim in range(3):
+        profile = np.ones(patch_size)
+
+        for i in range(border_size + fade_size):
+            if i < border_size:
+                weight = 0.1
+            else:
+                weight = 0.1 + 0.9 * (i - border_size) / fade_size
+
+            profile[i] = weight
+            profile[-(i + 1)] = weight
+
+        if dim == 0:
+            weights *= profile[:, np.newaxis, np.newaxis]
+        elif dim == 1:
+            weights *= profile[np.newaxis, :, np.newaxis]
+        else:
+            weights *= profile[np.newaxis, np.newaxis, :]
+
+    return weights
+
+
+def _segment_tile_list(tiles, model, batch_size=8, max_chunk_size=MAX_CHUNK_SIZE):
+    num_tiles = len(tiles)
+    segmented_tiles = []
+    for i in range(0, num_tiles, max_chunk_size):
+        chunk_end = min(i + max_chunk_size, num_tiles)
+        chunk = tiles[i:chunk_end]
+
+        try:
+            chunk_result = model.predict(chunk, verbose=0, batch_size=batch_size)
+            segmented_tiles.extend(chunk_result)
+
+        except tf.errors.ResourceExhaustedError:
+            print(f"Memory error with chunk size {len(chunk)}, falling back to smaller chunks")
+            fallback_chunk_size = max(1, len(chunk) // 4)
+            for j in range(i, chunk_end, fallback_chunk_size):
+                small_chunk = tiles[j:min(j + fallback_chunk_size, chunk_end)]
+                small_result = model.predict(small_chunk, verbose=0, batch_size=batch_size)
+                segmented_tiles.extend(small_result)
+
+    return segmented_tiles
+
+
+def _segment_tomogram_instance(volume, model, batch_size):
+    tiles, positions, original_shape = tile_volume(volume)
+    segmented_tiles = _segment_tile_list(tiles, model, batch_size=batch_size, max_chunk_size=MAX_CHUNK_SIZE)
+    segmented_tiles = np.array(segmented_tiles)
+    segmented_volume = detile_volume(segmented_tiles, positions, original_shape)
+
+    tf.keras.backend.clear_session()
+    gc.collect()
+    return segmented_volume.astype(np.float32)
+
+
+def segment_tomogram(model, tomogram_path, tta=1, batch_size=2):
+    volume = mrcfile.read(tomogram_path)
+    volume -= np.mean(volume)
+    volume /= np.std(volume) + 1e-8
+
+    segmented_volume = np.zeros_like(volume)
+    k_xy = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
+    k_fx = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
+    k_yz = [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1]
+    for j in range(tta):
+        tta_vol = volume.copy()
+        tta_vol = np.rot90(tta_vol, k=k_xy[j], axes=(1, 2))
+        tta_vol = tta_vol if not k_fx[j] else np.flip(tta_vol, axis=1)
+        tta_vol = np.rot90(tta_vol, k=2 * k_yz[j], axes=(0, 1))
+        segmented_tta_vol = _segment_tomogram_instance(tta_vol, model, batch_size)
+        segmented_tta_vol = np.rot90(segmented_tta_vol, k=-2 * k_yz[j], axes=(0, 1))
+        segmented_tta_vol = segmented_tta_vol if not k_fx[j] else np.flip(segmented_tta_vol, axis=1)
+        segmented_tta_vol = np.rot90(segmented_tta_vol, k=-k_xy[j], axes=(1, 2))
+        segmented_volume += segmented_tta_vol
+    segmented_volume /= tta
+
+    segmented_volume[:32, :, :] = 0
+    segmented_volume[-32:, :, :] = 0
+    segmented_volume[:, :32, :] = 0
+    segmented_volume[:, -32:, :] = 0
+    segmented_volume[:, :, :32] = 0
+    segmented_volume[:, :, -32:] = 0
+    return segmented_volume
+
+
+def save_mrc(pxd, path, data_format, voxel_size=10.0):
+    if data_format == 'float32':
+        pxd = pxd.astype(np.float32)
+    elif data_format == 'uint16':
+        pxd = (pxd * 255).astype(np.uint16) # scaling to [0, 255] because that's what we're used to in Ais
+    elif data_format == 'int8':
+        pxd = (pxd * 127).astype(np.int8)
+    with mrcfile.new(path, overwrite=True) as m:
+        m.set_data(pxd)
+        m.voxel_size = voxel_size
+
+def segmentation_thread(tomogram_list, model_path, feature, output_dir, gpu, batch_size, tta, overwrite, data_format):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    os.environ['TF_DISABLE_MKL'] = '1'
+    os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices=false'
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+    os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
+    for device in tf.config.list_physical_devices('GPU'):
+        tf.config.experimental.set_memory_growth(device, True)
+    mixed_precision.set_global_policy('mixed_float16')
+
+    process_start_time = psutil.Process().create_time()
+
+    model = load_model(model_path)
+
+    for j, tomogram_path in enumerate(tomogram_list, 1):
+        tomo_name = os.path.splitext(os.path.basename(tomogram_path))[0]
+        output_file = os.path.join(output_dir, f"{tomo_name}__{feature}.mrc")
+        wrote_temporary = False
+        try:
+            if os.path.exists(output_file):
+                file_age = os.path.getmtime(output_file)
+                if not overwrite or file_age > process_start_time - 60:
+                    continue
+
+            with mrcfile.new(output_file, overwrite=True) as m:
+                m.set_data(-1.0 * np.ones((10, 10, 10), dtype=np.float32))
+                wrote_temporary = True
+            segmented_volume = segment_tomogram(model, tomogram_path, tta, batch_size)
+
+            save_mrc(segmented_volume, output_file, data_format)
+            print(f"{j}/{len(tomogram_list)} (on GPU {gpu}) - {feature} - {os.path.basename(output_file)}")
+        except Exception as e:
+            if wrote_temporary:
+                os.remove(output_file)
+            print(f"{j}/{len(tomogram_list)} (on GPU {gpu}) - {feature} - {os.path.basename(output_file)} - ERROR: {e}")
+
+def dispatch_segment(feature, data_directory, output_directory, tta=1, gpus=(0,1,2,3), batch_size=8, overwrite=False, data_format='int8'):
+    if output_directory is None:
+        output_directory = data_directory
+    gpus = [int(n) for n in gpus.split(',')]
+
+    print(f'easymode segment\n'
+          f'feature: {feature}\n'
+          f'data_directory: {data_directory}\n'
+          f'output_directory: {output_directory}\n'
+          f'output_format: {data_format}\n'
+          f'gpus: {gpus}\n'
+          f'tta: {tta}\n'
+          f'overwrite: {overwrite}\n'
+          f'batch_size: {batch_size}\n')
+
+    tomograms = sorted(glob.glob(os.path.join(data_directory, '*.mrc')))
+
+    model_path = cache_model(feature)
+
+    os.makedirs(output_directory, exist_ok=True)
+
+    processes = []
+    for gpu in gpus:
+        p = multiprocessing.Process(target=segmentation_thread,
+                                    args=(tomograms, model_path, feature, output_directory, gpu, batch_size, tta, overwrite, data_format))
+        processes.append(p)
+        p.start()
+        time.sleep(2)
+
+    for p in processes:
+        p.join()
+
+
+
+
+
+