Upload 17 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Streamlit_app.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,79 @@
+# Unsupervised Segmentation App with Streamlit and PyTorch
+
+## Table of Contents
+1. [Introduction](#introduction)
+2. [Acknowledgments](#acknowledgments)
+3. [Requirements](#requirements)
+4. [Installation](#installation)
+5. [How to Run](#how-to-run)
+6. [Code Explanation](#code-explanation)
+7. [Contributing](#contributing)
+8. [License](#license)
+
+---
+
+## Introduction 🌟
+This project is a web application built using Streamlit and PyTorch. It performs unsupervised segmentation on uploaded images. The segmented image can be downloaded, and the colors of the segments can be customized.
+
+---
+
+## Acknowledgments 🙏
+This code is inspired from the project [pytorch-unsupervised-segmentation](https://github.com/kanezaki/pytorch-unsupervised-segmentation) by kanezaki. The original project is based on the paper "Unsupervised Image Segmentation by Backpropagation" presented at IEEE ICASSP 2018. The code is optimized for thin section images and microscopy analysis.
+
+---
+
+## Requirements 📋
+- Python 3.x
+- Streamlit
+- PyTorch
+- OpenCV
+- NumPy
+- scikit-image
+- PIL
+- base64
+
+---
+
+## Installation 🛠️
+
+1. **Clone the repository**
+    ```bash
+    git clone https://github.com/your-repo/unsupervised-segmentation.git
+    ```
+2. **Navigate to the project directory**
+    ```bash
+    cd unsupervised-segmentation
+    ```
+3. **Install the required packages**
+    ```bash
+    pip install -r requirements.txt
+    ```
+
 ---
-title: PetroSeg
-emoji: 🐨
-colorFrom: blue
-colorTo: blue
-sdk: streamlit
-sdk_version: 1.33.0
-app_file: app.py
-pinned: false
-license: mit
+
+## How to Run 🚀
+
+1. **Navigate to the project directory**
+    ```bash
+    cd unsupervised-segmentation
+    ```
+2. **Run the Streamlit app**
+    ```bash
+    streamlit run app.py
+    ```
+
 ---
+![Streamlit App Screenshot](https://github.com/fazzam12345/Unsupervised-Segmentation-App/blob/master/Streamlit_app.png?raw=true)
+
+
+---
+
+## Contributing 🤝
+Feel free to open issues and pull requests!
+
+---
+
+## License 📜
+This project is licensed under the MIT License.
+
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,6 @@
+from src import init
+from src.interface import main
+
+if __name__ == "__main__":
+    init.initialize_session_state()
+    main()

requirements.txt

@@ -0,0 +1,6 @@
+streamlit==1.27.0 
+opencv-python-headless
+numpy==1.24
+torch==2.0
+Pillow==9.5 
+scikit-image==0.18 

src/init.py

@@ -0,0 +1,5 @@
+import streamlit as st
+
+def initialize_session_state():
+    st.session_state.setdefault('segmented_image', None)
+    st.session_state.setdefault('new_colors', {})

src/interface.py

@@ -0,0 +1,139 @@
+
+import streamlit as st
+import cv2
+import numpy as np
+from src.models import perform_custom_segmentation
+from src.utils import resize_image, download_image
+import os
+import torch
+
+# Constants
+TARGET_SIZE = (750, 750)
+
+def get_parameters_from_sidebar() -> dict:
+    """Get segmentation parameters from sidebar"""
+    st.sidebar.header("Segmentation Parameters")
+    param_names = ['train_epoch', 'mod_dim1', 'mod_dim2', 'min_label_num', 'max_label_num']
+    param_values = [(1, 200, 43), (1, 128, 67), (1, 128, 63), (1, 20, 3), (1, 200, 25)]
+    params = {name: st.sidebar.slider(name.replace('_', ' ').title(), *values) for name, values in zip(param_names, param_values)}
+    
+    # Add sliders for target size width and height
+    target_size_width = st.sidebar.number_input("Target Size Width", 100, 1200, 750)
+    target_size_height = st.sidebar.number_input("Target Size Height", 100, 1200, 750)
+    params['target_size'] = (target_size_width, target_size_height)
+    
+    return params
+def display_segmentation_results() -> None:
+    """Display segmentation results"""
+    st.image(st.session_state.segmented_image, caption='Updated Segmented Image', use_column_width=True)
+
+def randomize_colors() -> None:
+    """Randomize colors for segmentation labels"""
+    unique_labels = np.unique(st.session_state.segmented_image.reshape(-1, 3), axis=0)
+    random_colors = {tuple(label): tuple(np.random.randint(0, 256, size=3)) for label in unique_labels}
+
+    for old_color, new_color in random_colors.items():
+        mask = np.all(st.session_state.segmented_image == np.array(old_color), axis=-1)
+        st.session_state.segmented_image[mask] = new_color
+
+    # Update color mappings in session state
+    st.session_state.new_colors.update(random_colors)
+    st.session_state.image_update_trigger += 1  # Trigger image update
+
+def handle_color_picking() -> None:
+    """Handle color picking and other functionalities"""
+    unique_labels = np.unique(st.session_state.segmented_image.reshape(-1, 3), axis=0)
+    for i, label in enumerate(unique_labels):
+        hex_label = f'#{label[0]:02x}{label[1]:02x}{label[2]:02x}'
+        new_color = st.color_picker(f"Choose a new color for label {i}", value=hex_label, key=f"label_{i}")
+        new_color_rgb = tuple(int(new_color.lstrip('#')[j:j+2], 16) for j in (0, 2, 4))
+        st.session_state.new_colors[tuple(label)] = new_color_rgb
+
+    # Convert the new colors to hexadecimal for comparison
+    new_colors_hex = {tuple(label): f'#{label[0]:02x}{label[1]:02x}{label[2]:02x}' for label in st.session_state.new_colors.values()}
+
+    for old_color, new_color in st.session_state.new_colors.items():
+        # Convert the old color to hexadecimal for comparison
+        old_color_hex = f'#{old_color[0]:02x}{old_color[1]:02x}{old_color[2]:02x}'
+        # Find the corresponding new color in hexadecimal
+        new_color_hex = new_colors_hex[new_color]
+        # Update the segmented image with the new color
+        mask = np.all(st.session_state.segmented_image == np.array(old_color), axis=-1)
+        st.session_state.segmented_image[mask] = new_color
+
+    # After updating colors, trigger an update to the segmented image display
+    st.session_state.image_update_trigger += 1
+
+def calculate_and_display_label_percentages() -> None:
+    """Calculate and display label percentages"""
+    final_labels = cv2.cvtColor(st.session_state.segmented_image, cv2.COLOR_BGR2GRAY)
+    unique_labels, counts = np.unique(final_labels, return_counts=True)
+    total_pixels = np.sum(counts)
+    label_percentages = {int(label): (count / total_pixels) * 100 for label, count in zip(unique_labels, counts)}
+
+    # Create a mapping from grayscale label to RGB color
+    label_to_color = {}
+    for label in unique_labels:
+        mask = final_labels == label
+        corresponding_color = st.session_state.segmented_image[mask][0]
+        hex_color = f'#{corresponding_color[0]:02x}{corresponding_color[1]:02x}{corresponding_color[2]:02x}'
+        label_to_color[int(label)] = hex_color
+
+    st.write("Label Percentages:")
+    for label, percentage in label_percentages.items():
+        hex_color = label_to_color[label]
+        color_box = f'<div style="display: inline-block; width: 20px; height: 20px; background-color: {hex_color}; margin-right: 10px;"></div>'
+        st.markdown(f'{color_box} Label {label}: {percentage:.2f}%', unsafe_allow_html=True)
+
+def main() -> None:
+    st.title("PetroSeg")
+    st.info("""
+    - **Training Epochs**: Higher values will lead to fewer segments but may take more time.
+    - **Image Size**: For better efficiency, upload small-sized images.
+    - **Cache**: For best results, clear the cache between different image uploads. You can do this from the menu in the top-right corner.
+    """)
+
+    if torch.cuda.is_available():
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
+
+    # Initialize session state if not already initialized
+    if 'segmented_image' not in st.session_state:
+        st.session_state.segmented_image = None
+    if 'new_colors' not in st.session_state:
+        st.session_state.new_colors = {}
+    if 'image_update_trigger' not in st.session_state:
+        st.session_state.image_update_trigger = 0
+
+    # Define params before using it
+    params = get_parameters_from_sidebar()
+
+    uploaded_image = st.sidebar.file_uploader("Upload an image", type=["jpg", "png", "jpeg", "bmp", "tiff", "webp"])
+    if uploaded_image:
+        file_bytes = np.asarray(bytearray(uploaded_image.read()), dtype=np.uint8)
+        image = cv2.imdecode(file_bytes, 1)
+
+        if image is None:
+            st.error("Error loading image. Please check the file and try again.")
+            return
+
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        st.image(image_rgb, caption='Original Image', use_column_width=True)
+
+        # Use the target size specified by the user
+        target_size = params['target_size']
+        image_resized = resize_image(image_rgb, target_size)
+
+        if st.sidebar.button("Start Segmentation"):
+            st.session_state.segmented_image = perform_custom_segmentation(image_resized, params)
+
+        if st.sidebar.button("Change Colors"):
+            randomize_colors()
+
+        if st.session_state.segmented_image is not None:
+            handle_color_picking()
+            display_segmentation_results()
+            calculate_and_display_label_percentages()
+            download_image(st.session_state.segmented_image, 'segmented_image.png')
+
+if __name__ == "__main__":
+    main()

src/models.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import streamlit as st
+import os
+from skimage import segmentation
+
+
+def perform_custom_segmentation(image, params):
+    class Args(object):
+        def __init__(self, params):
+            self.train_epoch = params.get('train_epoch', 2 ** 3)
+            self.mod_dim1 = params.get('mod_dim1', 64)
+            self.mod_dim2 = params.get('mod_dim2', 32)
+            self.gpu_id = params.get('gpu_id', 0)
+            self.min_label_num = params.get('min_label_num', 6)
+            self.max_label_num = params.get('max_label_num', 256)
+
+    args = Args(params)
+
+    class MyNet(nn.Module):
+        def __init__(self, inp_dim, mod_dim1, mod_dim2):
+            super(MyNet, self).__init__()
+            self.seq = nn.Sequential(
+                nn.Conv2d(inp_dim, mod_dim1, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(mod_dim1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(mod_dim1, mod_dim2, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(mod_dim2),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(mod_dim2, mod_dim1, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(mod_dim1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(mod_dim1, mod_dim2, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(mod_dim2),
+            )
+
+        def forward(self, x):
+            return self.seq(x)
+
+    torch.cuda.manual_seed_all(1943)
+    np.random.seed(1943)
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
+
+    '''segmentation ML'''
+    seg_map = segmentation.felzenszwalb(image, scale=15, sigma=0.06, min_size=14)
+    seg_map = seg_map.flatten()
+    seg_lab = [np.where(seg_map == u_label)[0]
+               for u_label in np.unique(seg_map)]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
+    tensor = image.transpose((2, 0, 1))
+    tensor = tensor.astype(np.float32) / 255.0
+    tensor = tensor[np.newaxis, :, :, :]
+    tensor = torch.from_numpy(tensor).to(device)
+
+    model = MyNet(inp_dim=3, mod_dim1=args.mod_dim1, mod_dim2=args.mod_dim2).to(device)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=5e-2, momentum=0.9)
+
+    image_flatten = image.reshape((-1, 3))
+    color_avg = np.random.randint(255, size=(args.max_label_num, 3))
+    show = image
+
+    progress_bar = st.progress(0)
+
+    for batch_idx in range(args.train_epoch):
+        optimizer.zero_grad()
+        output = model(tensor)[0]
+        output = output.permute(1, 2, 0).view(-1, args.mod_dim2)
+        target = torch.argmax(output, 1)
+        im_target = target.data.cpu().numpy()
+
+        for inds in seg_lab:
+            u_labels, hist = np.unique(im_target[inds], return_counts=True)
+            im_target[inds] = u_labels[np.argmax(hist)]
+
+        target = torch.from_numpy(im_target)
+        target = target.to(device)
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
+        un_label, lab_inverse = np.unique(im_target, return_inverse=True, )
+        if un_label.shape[0] < args.max_label_num:
+            img_flatten = image_flatten.copy()
+            if len(color_avg) != un_label.shape[0]:
+                color_avg = [np.mean(img_flatten[im_target == label], axis=0, dtype=int) for label in un_label]
+            for lab_id, color in enumerate(color_avg):
+                img_flatten[lab_inverse == lab_id] = color
+            show = img_flatten.reshape(image.shape)
+
+        progress = (batch_idx + 1) / args.train_epoch
+        progress_bar.progress(progress)
+
+    return show

src/utils.py

@@ -0,0 +1,39 @@
+import cv2
+import tempfile
+from io import BytesIO
+import streamlit as st
+import numpy as np
+
+def resize_image(image, size):
+    return cv2.resize(image, size, interpolation=cv2.INTER_AREA)
+
+def automatically_change_segment_colors(segmented_image):
+    # Generate a unique color for each segment
+    unique_labels = np.unique(segmented_image.reshape(-1, 3), axis=0)
+    new_colors = np.random.randint(0, 256, (len(unique_labels), 3), dtype=np.uint8)
+    
+    # Apply the new colors to the segmented image
+    for i, label in enumerate(unique_labels):
+        mask = np.all(segmented_image == label, axis=-1)
+        segmented_image[mask] = new_colors[i]
+    
+    return segmented_image
+
+def download_image(image_array, file_name):
+    try:
+        image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+        success = cv2.imwrite(temp_file.name, image_array)
+        if not success:
+            st.error("Could not save image.")
+            return
+        with open(temp_file.name, 'rb') as f:
+            bytes = f.read()
+        st.download_button(
+            label="Download Image",
+            data=BytesIO(bytes),
+            file_name=file_name,
+            mime='image/png',
+        )
+    except Exception as e:
+        st.error(f"An error occurred: {e}")