update code

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

.gitignore

@@ -0,0 +1,4 @@
+# Ignore Python cache files
+*.pyc
+__pycache__/
+venv/

.python-version

@@ -0,0 +1 @@
+3.11.0

ML_model/data/X_train.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8104f7e7cbc69e960b07138448d538edc2d6856d02862a4182466da58ec07ab7
+size 2775817

ML_model/data/selected_features.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8593b584c411a5f65c5a5265e987f077d54f8e17fbe7ae9b653a69d93b83955b
+size 344

ML_model/model.py

@@ -0,0 +1,66 @@
+# from sklearn.ensemble import RandomForestClassifier
+import numpy as np
+import joblib
+import pandas as pd
+from sklearn.preprocessing import MinMaxScaler
+from lime import lime_tabular
+
+class MLModel():
+    def __init__(self):
+        self.model = joblib.load('./ML_model/weights/xgboost_convnet_best.pkl')
+
+        self.clinical_data  = pd.read_csv("./ML_model/data/clinical-2.csv")
+        del self.clinical_data['ID']
+        del self.clinical_data['SIDE']
+
+        self.columns_to_scale = ['AGE', 'HEIGHT', 'WEIGHT', 'MAX WEIGHT', 'BMI', 'KOOS PAIN SCORE']
+        self.scaler = MinMaxScaler()
+        self.clinical_data[self.columns_to_scale] = self.scaler.fit_transform(self.clinical_data[self.columns_to_scale])
+
+        self.columns_to_convert = ['FREQUENT PAIN', 'SURGERY', 'RISK', 'SXKOA', 'SWELLING', 'BENDING FULLY', 'SYMPTOMATIC', 'CREPITUS']
+        self.mapping_dict = {}
+        for column in self.columns_to_convert:
+            self.clinical_data[column], unique_values = pd.factorize(self.clinical_data[column])
+            self.mapping_dict[column] = unique_values
+
+        self.selected_features = joblib.load("./ML_model/data/selected_features.pkl")
+        self.X_train = joblib.load("./ML_model/data/X_train.pkl")
+        self.selected_feature_index = [self.X_train.columns.get_loc(col) for col in self.selected_features]
+
+        self.explainer = lime_tabular.LimeTabularExplainer(
+            self.X_train[self.selected_features].to_numpy(),
+            feature_names=self.selected_features,
+            class_names=['0', '1', '2', '3', '4'],
+            mode='classification'
+        )
+                
+    def get_clinical_data(self, filename):
+        row = self.clinical_data.loc[self.clinical_data['FILENAME'] == filename]
+        return row.to_numpy()# remove filename col
+
+    def predict(self, overal_diagnosis, jsws, clinical=None, filename=None):
+        if not clinical:
+            if not filename:
+                raise "Need clinical data or filename in OAI database"
+            clinical = self.get_clinical_data(filename)
+            assert clinical.shape[0] != 0, "clinical data not found"
+        else:
+            raise "not implemented yet!"
+
+        x = np.concatenate((overal_diagnosis, clinical[0, 1:], np.array(jsws)))
+
+        return self.model.predict(x = x[self.selected_feature_index].reshape(1, -1)) # unsqueeze(0) 
+    
+    def predict_explain(self, overal_diagnosis, jsws, clinical=None, filename=None):
+        if not clinical:
+            if not filename:
+                raise "Need clinical data or filename in OAI database"
+            clinical = self.get_clinical_data(filename)
+            assert clinical.shape[0] != 0, "clinical data not found"
+        else:
+            raise "not implemented yet!"
+        
+        x = np.concatenate((overal_diagnosis, clinical[0, 1:], np.array(jsws)))[self.selected_feature_index]
+
+        exp = self.explainer.explain_instance(x, self.model.predict_proba, num_features=len(self.selected_features), top_labels=1)
+        return self.model.predict(x.reshape(1, -1))[0], exp

ML_model/weights/xgboost_convnet_best.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ee491112dc64c61cc66b9922262c1a5920db5270dc6cf1220355a737b1c30c0
+size 286017

README.md

@@ -0,0 +1,12 @@
+---
+title: Demo XDesco
+emoji: 💻
+colorFrom: red
+colorTo: yellow
+sdk: streamlit
+sdk_version: 1.32.2
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,192 @@
+import streamlit as st
+import cv2
+import matplotlib.pyplot as plt
+import matplotlib.patheffects as path_effects
+from PIL import Image
+import numpy as np
+import os
+import streamlit.components.v1 as components
+import pandas as pd 
+
+from utils import read_image, combine_mask, combine_prob_jsw, scale_coordinates, get_annotations
+from segmentation.model import Segmenter
+from jsw import get_JSW, calculate_diff, calculate_jsw_info
+from classification.model import Classifier
+from ML_model.model import MLModel
+from gae import AnomalyExtractor
+
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+
+st.set_page_config(layout="wide")
+
+@st.cache_resource
+def load_models():
+    seg_model = Segmenter()
+    classif_model = Classifier('convnet')
+    ml_model = MLModel()
+    anomaly_extractor = AnomalyExtractor()
+
+    return seg_model, classif_model, ml_model, anomaly_extractor
+
+
+seg_model, classif_model, ml_model, anomaly_extractor = load_models()
+
+st.sidebar.title("Demo xDesco System")
+
+uploaded_file = st.sidebar.file_uploader("Choose a knee join x-ray image", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    file_bytes = uploaded_file.read()
+    file_name = uploaded_file.name
+    img = read_image(file_bytes)
+
+    mask = seg_model.segment(img)
+
+    mask_image = combine_mask(img, mask)
+    col1, col2 = st.columns(2)
+
+
+
+    left_distances, right_distances, links = get_JSW(mask, dim = 10, verbose = 0)
+    left_links, right_links = links
+    jsw_m, jsw_mm = calculate_diff(left_distances, right_distances)
+    diff_percentage, mean_left, mean_right, side_min, index_min, value_min = calculate_jsw_info(left_distances, right_distances)
+
+    scaled_left_links = scale_coordinates([coord for pair in left_links for coord in pair], original_size=224, mask_size=640)
+    left_pairs = list(zip(scaled_left_links[::2], scaled_left_links[1::2]))
+    scaled_right_links = scale_coordinates([coord for pair in right_links for coord in pair], original_size=224, mask_size=640)
+    right_pairs = list(zip(scaled_right_links[::2], scaled_right_links[1::2]))
+
+    # st.write("Left JSW: [{}]".format(", ".join(f"{x:.2f}" for x in left_distances)))
+    # st.write("Right JSW: [{}]".format(", ".join(f"{x:.2f}" for x in right_distances)))
+
+    # st.write("$JSW_M$: ", jsw_m)
+    # st.write("$JSW_{MM}$: ", jsw_mm)
+
+    probabilites = classif_model.predict(img)
+    # st.write("Probability: [{}]".format(", ".join(f"{x:.2f}" for x in probabilites[0])))
+    # st.write("len: ", len(probabilites[0]))
+    annotations = get_annotations(probabilites[0])
+
+    predicted, exp = ml_model.predict_explain(probabilites[0], [jsw_m, jsw_mm], filename=file_name)
+
+    processed_img, anomaly = anomaly_extractor.extract(mask, img, verbose=0)
+
+    def plot_anomaly_with_clues(processed_img, anomaly, diff_percentage, mean_left, mean_right, side_min, index_min, value_min,
+                                left_pairs, right_pairs, color='r', thickness=1):
+        # Tạo một figure và axes từ matplotlib
+        fig, ax = plt.subplots()
+        ax.imshow(processed_img, cmap = 'gray')
+        ax.imshow(anomaly, cmap='turbo', alpha = 0.3)  
+
+        # # Vẽ đường thẳng từ pairs_left
+        # for pairs in [left_pairs, right_pairs]:
+        #     for pair in pairs:
+        #         start_point = pair[0]
+        #         end_point = pair[1]
+        #         ax.plot([start_point[0], end_point[0]], [start_point[1], end_point[1]], color=color, linewidth=thickness)
+        #         ax.scatter(*start_point, c=color, s=thickness*2)  # Vẽ điểm đầu
+        #         ax.scatter(*end_point, c=color, s=thickness*2)    # Vẽ điểm cuối
+        
+        # Thêm chữ diff_percentage vào giữa ảnh
+        mid_x = anomaly.shape[1] // 2
+        mid_y = anomaly.shape[0] // 10
+        # ax.text(mid_x, mid_y, f'% difference between left & right joint space distance: {diff_percentage:.2f}%', 
+        #         color='white', fontsize=8, ha='center', va='center')
+
+
+        # Thêm mean_distance cho left và right
+        # left_mid_index = len(left_pairs) // 2
+        # right_mid_index = len(right_pairs) // 2
+        # left_mid_point = left_pairs[left_mid_index][0]
+        # right_mid_point = right_pairs[right_mid_index][0]
+
+        # ax.text(left_mid_point[0], anomaly.shape[0] // 1.5, f'{mean_left:.2f}', color='yellow', fontsize=20, ha='center', va='center', path_effects=[path_effects.Stroke(linewidth=2, foreground='black'), path_effects.Normal()])
+        # ax.text(right_mid_point[0], anomaly.shape[0] // 1.5, f'{mean_right:.2f}', color='yellow', fontsize=20, ha='center', va='center', path_effects=[path_effects.Stroke(linewidth=2, foreground='black'), path_effects.Normal()])
+        # print((mean_left // value_min > 2) if side_min == 0 else (mean_right // value_min > 2))
+        
+        if (diff_percentage > 0) or ((mean_left / value_min > 2) if side_min == 0 else (mean_right / value_min > 2)):
+            min_pairs = left_pairs if side_min == 0 else right_pairs
+            min_pair = min_pairs[index_min]
+            start_point = min_pair[0]
+            end_point = min_pair[1]
+
+            # Xác định tọa độ để vẽ bbox xung quanh đường thẳng đứng với padding
+            padding_x = 23  # Độ rộng padding theo chiều x
+            padding_y = 13  # Độ rộng padding theo chiều y
+            min_x = start_point[0] - padding_x
+            max_x = start_point[0] + padding_x
+            min_y = min(start_point[1], end_point[1]) - padding_y
+            max_y = max(start_point[1], end_point[1]) + padding_y
+
+            rect = plt.Rectangle((min_x, min_y), max_x - min_x, max_y - min_y, linewidth=3, edgecolor='red', facecolor='none')
+            ax.add_patch(rect)
+
+            # Thêm văn bản với giá trị value_min tại vị trí của đường thẳng có khoảng cách nhỏ nhất
+            # text = ax.text(start_point[0], min_y-padding_y//2, f'{value_min:.2f}', color='red', fontsize=18, ha='center', va='center', path_effects=[path_effects.Stroke(linewidth=1.5, foreground='black'), path_effects.Normal()])
+            
+        # add annotation osteophyte & jsn
+        ax.text(anomaly.shape[1] // 2, anomaly.shape[0] // 3.2, f'{annotations["osteophyte"]}', color='yellow', fontsize=19, ha='center', va='center', path_effects=[path_effects.Stroke(linewidth=1.5, foreground='black'), path_effects.Normal()])
+        ax.text(anomaly.shape[1] // 2, anomaly.shape[0] // 1.3, f'{annotations["jsn"]}', color='yellow', fontsize=19, ha='center', va='center', path_effects=[path_effects.Stroke(linewidth=1.5, foreground='black'), path_effects.Normal()])
+
+        # add annoatation jsw_m & jsw_mm info
+        ax.text(anomaly.shape[1] // 40, anomaly.shape[0] // 1.05, f'$JSW_{{Mean}}$: {jsw_m:.2f}', color='white', fontsize=10, ha='left', va='center', path_effects=[path_effects.Stroke(linewidth=1, foreground='black'), path_effects.Normal()])
+        ax.text(anomaly.shape[1] // 40, anomaly.shape[0] // 1.1, f'$JSW_{{MM}}$: {jsw_mm:.2f}', color='white', fontsize=10, ha='left', va='center', path_effects=[path_effects.Stroke(linewidth=1, foreground='black'), path_effects.Normal()])
+
+        ax.axis('off')
+        # ax.set_title('Annotated anomaly map')
+        # plt.show()
+        return fig
+
+    with col1:
+        caption = "Uploaded Image"
+        st.markdown(
+            f"<h3 style='text-align: center;'>{caption}</h3>", 
+            unsafe_allow_html=True
+        )
+        st.image(img, channels="BGR", use_column_width=True)
+        # plt.imshow(anomaly, cmap="turbo")
+        # plt.axis('off')
+        # st.pyplot()
+    with col2:
+        # st.image(mask_image, channels="BGR", caption='mask image', use_column_width=True)
+
+        fig = plot_anomaly_with_clues(
+            processed_img,
+            anomaly,
+            diff_percentage,
+            mean_left, 
+            mean_right, 
+            side_min, index_min, 
+            value_min, 
+            left_pairs, 
+            right_pairs, 
+            color='r', 
+            thickness=1
+        )
+        caption = "Annotated Anomaly Map"
+        st.markdown(
+            f"<h3 style='text-align: center;'>{caption}</h3>", 
+            unsafe_allow_html=True
+        )
+        st.pyplot(fig, bbox_inches='tight', pad_inches=0)
+        
+
+    exp_html = exp.as_html()
+    components.html(exp_html,  height=1400, width=None)
+
+    # def extract_explain(exp, label):
+    #     ans = exp.local_exp[label]
+    #     ans = [(exp.domain_mapper.feature_names[x[0]],
+    #             exp.domain_mapper.feature_values[x[0]],
+    #             exp.domain_mapper.discretized_feature_names[x[0]],
+    #             float(x[1])
+    #             ) for x in ans]
+        
+    #     return ans
+
+    # explanation_list = extract_explain(exp, predicted)
+    # explanation_df = pd.DataFrame(explanation_list, columns=['Feature', 'Value', 'Explain', 'Weight'])
+    # st.table(explanation_df)
+

classification/model.py

@@ -0,0 +1,77 @@
+import torch
+import torchvision
+import cv2
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+
+from .utils import to_rgb
+
+
+def create_effNetv2s():
+    model = torchvision.models.efficientnet_v2_s(weights='IMAGENET1K_V1')
+    num_features = model.classifier[1].in_features
+    model.classifier[1] = nn.Sequential(
+                nn.Linear(num_features, NUM_CLASSES),
+                nn.Sigmoid()
+            )
+    return model
+
+def create_convnet():
+    model = torchvision.models.convnext_base(weights='IMAGENET1K_V1')
+    num_features = model.classifier[2].in_features
+    model.classifier[2] = nn.Sequential(
+                nn.Linear(num_features, NUM_CLASSES),
+                nn.Sigmoid()
+            )
+    return model
+
+def create_model(model_name):
+    model = _MODEL[model_name]()
+    model.load_state_dict(torch.load(_WEIGHT[model_name], map_location=torch.device('cpu')))
+    model.to(DEVICE)
+    return model
+
+def create_transform():
+    transform = transforms.Compose([
+        transforms.Resize((HEIGHT, WEIGHT)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.6078, 0.6078, 0.6078), (0.1932, 0.1932, 0.1932))
+    ])
+    return transform
+
+
+_MODEL = {
+    "effNetv2s": create_effNetv2s,
+    "convnet": create_convnet
+}
+
+_WEIGHT = {
+    "effNetv2s": './classification/weights/effnetv2s.pt',
+    "convnet": './classification/weights/convnet.pt',
+}
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+HEIGHT = 224
+WEIGHT = 224
+NUM_CLASSES = 44
+
+
+class Classifier():
+    def __init__(self, model_name="effNetv2s"):
+        self.model = create_model(model_name)
+        self.transform = create_transform()
+
+    def predict(self, image):
+        '''
+            input: cv2 image
+            output: multi-label probability vector
+        '''
+        image = to_rgb(image)
+        image = Image.fromarray(image)
+        image = self.transform(image)
+        self.model.eval()
+        with torch.no_grad():
+            out = self.model(image.unsqueeze(0).to(DEVICE)).cpu().numpy()
+            return out

classification/utils.py

@@ -0,0 +1,6 @@
+import cv2
+
+def to_rgb(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

classification/weights/convnet.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe61c8220bbd41af3e536451985c1c49cc0d8b438de2da76f7fdbba8f6051741
+size 350549706

gae/__init__.py

@@ -0,0 +1 @@
+from .anomaly_extraction import AnomalyExtractor

gae/anomaly_extraction.py

@@ -0,0 +1,107 @@
+import matplotlib.pyplot as plt
+import streamlit as st
+import tensorflow as tf
+from PIL import Image
+from skimage.filters import gaussian
+
+from .model import DCGAN
+from .modules import DropBlockNoise
+from .utils import *
+from .contour import get_contours_v2
+
+
+IMAGE_SIZE = 224
+
+
+class AnomalyExtractor():
+    def __init__(self):
+        print("================================load dcgan=================================\n")
+        dcgan = DCGAN(input_shape=(IMAGE_SIZE, IMAGE_SIZE, 1),
+              architecture='two-stage',
+              output_activation='sigmoid',
+              noise=DropBlockNoise(rate=0.1, block_size=16),
+              pretrain_weights=None,
+              block_type='pixel-shuffle',
+              kernel_initializer='glorot_uniform',
+              C=1.)
+        
+        self.restore_model = dcgan.generator
+        self.restore_model.load_weights("./gae/weights/gan_efficientunet_full_augment-hist_equal_generator.h5")
+        self.restore_model.trainable = False
+
+    def extract(self, mask, img, verbose=0):
+        img = to_gray(img)
+        img = np.array(Image.fromarray(img).resize((224, 224)))
+        img = preprocess(img)
+
+        if verbose:
+            # Hiển thị ảnh gốc
+            show_image(img, title="Original image")
+            plt.axis('off')
+            st.pyplot()
+
+            # Hiển thị mask
+            show_image(mask, title="Mask")
+            plt.axis('off')
+            st.pyplot()
+
+        uc, lc = get_contours_v2(mask, verbose=0)
+
+        mask = np.zeros((640, 640)).astype('uint8')
+        mask = draw_points(mask, lc, thickness=1, color=(255, 255, 255))
+        mask = draw_points(mask, uc, thickness=1, color=(255, 255, 255))
+        mask = cv2.resize(mask, (224, 224), cv2.INTER_NEAREST)
+        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+        mask = mask / 255.
+
+        if verbose:
+            show_image(mask, title = "Contour")
+            plt.axis('off')
+            st.pyplot()
+
+        # sử dụng equalization histogram
+        mask = 1 - mask
+        dila = dilate(mask)
+        dilated = gaussian(dilate(mask), sigma=50, truncate=0.3)
+
+        if verbose:
+            show_image(dila, title="dilated Image")
+            plt.axis('off')
+            st.pyplot()
+
+            show_image(dilated, title="blur Image")
+            plt.axis('off')
+            st.pyplot()
+
+        im = np.expand_dims(img * (1 - dilated), axis=0)
+        im = tf.convert_to_tensor(im, dtype=tf.float32)
+
+        restored_img = self.restore_model(im)
+
+        res = tf.squeeze(tf.squeeze(restored_img, axis=-1), axis=0)
+
+        if verbose:
+            show_image(im[0], title="Masked Image")
+            plt.axis('off')
+            st.pyplot()
+
+            show_image(res, title="Reconstructed image")
+            plt.axis('off')
+            st.pyplot()
+
+            show_image(dilated*tf.abs(img - res), title="Anomaly map", cmap_type='turbo')
+            plt.axis('off')
+            st.pyplot()
+
+
+        # plt.imshow(img, cmap = 'gray')
+        # plt.imshow(dilated*tf.abs(img - res), cmap ='turbo', alpha = 0.3)
+        # plt.axis('off')
+        # plt.show()
+        # st.pyplot()
+
+        # show_image(dilated, title="dilated", cmap_type='gray')
+        # plt.axis('off')
+        # st.pyplot()
+
+        return img, dilated*tf.abs(img - res)

gae/contour.py

@@ -0,0 +1,64 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+
+
+def find_boundaries_v2(mask, top=True, verbose=0):
+    boundaries = []
+    height, width = mask.shape
+
+    contours, _ = cv2.findContours(255 * mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+
+    areas = np.array([cv2.contourArea(cnt) for cnt in contours])
+    contour = contours[areas.argmax()]
+    contour = contour.reshape(-1, 2)
+    org_contour = contour.copy()
+    pos = (contour[:, 1].max() + contour[:, 1].min()) // 2
+    idx = np.where(contour[:, 1] == pos)
+    if contour[idx[0][0]][0] < contour[idx[0][-1]][0] and not top:
+        # start = contour[idx[0][0]]
+        # end = contour[idx[0][-1]]
+        start_idx = idx[0][0]
+        end_idx = idx[0][-1]
+    else:
+        # end = contour[idx[0][0]]
+        # start = contour[idx[0][-1]]
+        end_idx = idx[0][0]
+        start_idx = idx[0][-1]
+    # start_idx = ((start - contour) ** 2).sum(axis=-1).argmin()
+    # end_idx = ((end - contour) ** 2).sum(axis=-1).argmin()
+    if start_idx <= end_idx:
+        contour = contour[start_idx:end_idx + 1]
+    else:
+        contour = np.concatenate([contour[start_idx:], contour[:end_idx + 1]])
+    if verbose:
+        temp = draw_points(127 * mask.astype(np.uint8), contour, thickness=5)
+        temp = draw_points(temp, [start, end], color=[155, 155], thickness=15)
+        cv2_imshow(temp)
+
+    return np.array(contour), np.array(org_contour)
+
+
+def get_contours_v2(mask, verbose=0):
+    upper_contour, full_upper = find_boundaries_v2(mask == 1, top=False, verbose=verbose)
+    lower_contour, full_lower = find_boundaries_v2(mask == 2, top=True, verbose=verbose)
+    if verbose:
+        temp = draw_points(127 * mask, full_upper, thickness=3, color=(255, 0, 0))
+        temp = draw_points(temp, full_lower, thickness=3)
+        plt.imshow(temp)
+        plt.title("Segmentation")
+        plt.axis('off')
+        plt.show()
+        # st.pyplot()
+        # cv2.imwrite('full.png', temp)
+    #         temp = draw_points(temp, limit_points, thickness = 7, color = (0, 0, 255))
+    #         cv2_imshow(temp)
+    #         cv2.imwrite('limit_points.png', temp)
+    if verbose:
+        temp = draw_points(127 * mask, upper_contour, thickness=3, color=(255, 0, 0))
+        temp = draw_points(temp, lower_contour, thickness=3)
+        cv2_imshow(temp)
+        # st.pyplot()
+        # cv2.imwrite('cropped.png', temp)
+
+    return upper_contour, lower_contour

gae/model.py

@@ -0,0 +1,449 @@
+import gc
+import numpy as np
+import cv2
+from tqdm import tqdm
+
+import tensorflow as tf
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.models import Sequential, Model
+
+from tensorflow.keras import layers
+from tensorflow.keras.applications import EfficientNetV2S
+from tensorflow.keras.layers import (
+    Dense, Flatten, Conv2D, Activation, BatchNormalization,
+    MaxPooling2D, AveragePooling2D, GlobalAveragePooling2D,
+    Dropout, Input, concatenate, add, Conv2DTranspose, Lambda,
+    SpatialDropout2D, Cropping2D, UpSampling2D, LeakyReLU,
+    ZeroPadding2D, Reshape, Concatenate, Multiply, Permute, Add
+)
+
+from .contour import get_contours_v2
+from .modules import (
+    MultipleTrackers, DropBlockNoise, squeeze_excite_block, spatial_squeeze_excite_block,
+    channel_spatial_squeeze_excite, DoubleConv, UpSampling2D_block, Conv2DTranspose_block,
+    PixelShuffle_block
+)
+from .utils import mae
+
+
+IMAGE_SIZE = 224
+
+
+def adjust_pretrained_weights(model_cls, input_size, name=None):
+    weights_model = model_cls(weights='imagenet',
+                              include_top=False,
+                              input_shape=(*input_size, 3))
+    target_model = model_cls(weights=None,
+                             include_top=False,
+                             input_shape=(*input_size, 1))
+    weights = weights_model.get_weights()
+    weights[0] = np.sum(weights[0], axis=2, keepdims=True)
+    target_model.set_weights(weights)
+
+    del weights_model
+    tf.keras.backend.clear_session()
+    gc.collect()
+    if name:
+        target_model._name = name
+    return target_model
+
+
+def get_efficient_unet(name=None,
+                       option='full',
+                       input_shape=(IMAGE_SIZE, IMAGE_SIZE, 1),
+                       encoder_weights=None,
+                       block_type='conv-transpose',
+                       output_activation='sigmoid',
+                       kernel_initializer='glorot_uniform'):
+
+    if encoder_weights == 'imagenet':
+        encoder = adjust_pretrained_weights(EfficientNetV2S, input_shape[:-1], name)
+    elif encoder_weights is None:
+        encoder = EfficientNetV2S(weights=None,
+                                  include_top=False,
+                                  input_shape=input_shape)
+        encoder._name = name
+    else:
+        raise ValueError(encoder_weights)
+
+    if option == 'encoder':
+        return encoder
+
+    MBConvBlocks = []
+
+    skip_candidates = ['1b', '2d', '3d', '4f']
+
+    for mbblock_nr in skip_candidates:
+        mbblock = encoder.get_layer('block{}_add'.format(mbblock_nr)).output
+        MBConvBlocks.append(mbblock)
+
+    head = encoder.get_layer('top_activation').output
+    blocks = MBConvBlocks + [head]
+
+    if block_type == 'upsampling':
+        UpBlock = UpSampling2D_block
+    elif block_type == 'conv-transpose':
+        UpBlock = Conv2DTranspose_block
+    elif block_type == 'pixel-shuffle':
+        UpBlock = PixelShuffle_block
+    else:
+        raise ValueError(block_type)
+
+    o = blocks.pop()
+    o = UpBlock(512, initializer=kernel_initializer, skip=blocks.pop())(o)
+    o = UpBlock(256, initializer=kernel_initializer, skip=blocks.pop())(o)
+    o = UpBlock(128, initializer=kernel_initializer, skip=blocks.pop())(o)
+    o = UpBlock(64, initializer=kernel_initializer, skip=blocks.pop())(o)
+    o = UpBlock(32, initializer=kernel_initializer, skip=None)(o)
+    o = Conv2D(input_shape[-1], (1, 1), padding='same', activation=output_activation, kernel_initializer=kernel_initializer)(o)
+
+    model = Model(encoder.input, o, name=name)
+
+    if option == 'full':
+        return model, encoder
+    elif option == 'model':
+        return model
+    else:
+        raise ValueError(option)
+
+
+class DCGAN():
+    def __init__(self,
+                 input_shape=(IMAGE_SIZE, IMAGE_SIZE, 1),
+                 architecture='two-stage',
+                 pretrain_weights=None,
+                 output_activation='sigmoid',
+                 block_type='conv-transpose',
+                 kernel_initializer='glorot_uniform',
+                 noise=None,
+                 C=1.):
+
+        self.C = C
+        # Build
+        kwargs = dict(input_shape=input_shape,
+                      output_activation=output_activation,
+                      encoder_weights=pretrain_weights,
+                      block_type=block_type,
+                      kernel_initializer=kernel_initializer)
+
+        if architecture == 'two-stage':
+            encoder = get_efficient_unet(name='dcgan_disc',
+                                         option='encoder',
+                                         **kwargs)
+
+            self.generator = get_efficient_unet(name='dcgan_gen', option='model', **kwargs)
+        elif architecture == 'shared':
+
+            self.generator, encoder = get_efficient_unet(name='dcgan', option='full', **kwargs)
+        else:
+            raise ValueError(f'Unsupport architecture: {architecture}')
+
+        gpooling = GlobalAveragePooling2D()(encoder.output)
+        prediction = Dense(1, activation='sigmoid')(gpooling)
+        self.discriminator = Model(encoder.input, prediction, name='dcgan_disc')
+
+        tf.keras.backend.clear_session()
+        _ = gc.collect()
+
+        if noise:
+            gen_inputs = self.generator.input
+            corrupted_inputs = noise(gen_inputs)
+            outputs = self.generator(corrupted_inputs)
+            self.generator = Model(gen_inputs, outputs, name='dcgan_gen')
+
+            tf.keras.backend.clear_session()
+            _ = gc.collect()
+
+        if output_activation == 'tanh':
+
+            self.process_input = layers.Lambda(lambda img: (img*2.-1.), name='dcgan_normalize')
+            self.process_output = layers.Lambda(lambda img: (img*0.5+0.5), name='dcgan_denormalize')
+            gen_inputs = self.generator.input
+            process_inputs = self.process_input(gen_inputs)
+            process_inputs = self.generator(process_inputs)
+            gen_outputs = self.process_output(process_inputs)
+            self.generator = Model(gen_inputs, gen_outputs, name='dcgan_gen')
+
+            disc_inputs = self.discriminator.input
+            process_inputs = self.process_input(disc_inputs)
+            disc_outputs = self.discriminator(process_inputs)
+            self.discriminator = Model(disc_inputs, disc_outputs, name='dcgan_disc')
+
+            tf.keras.backend.clear_session()
+            _ = gc.collect()
+
+    def summary(self):
+        self.generator.summary()
+        self.discriminator.summary()
+
+    def compile(self,
+                generator_optimizer=Adam(5e-4, 0.5),
+                discriminator_optimizer=Adam(5e-4),
+                reconstruction_loss=mae,
+                discriminative_loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
+                reconstruction_metrics=[],
+                discriminative_metrics=[]):
+
+        self.discriminator_optimizer = discriminator_optimizer
+        self.discriminator.compile(optimizer=self.discriminator_optimizer)
+
+        self.generator_optimizer = generator_optimizer
+        self.generator.compile(optimizer=self.generator_optimizer)
+
+        self.loss = discriminative_loss
+        self.reconstruction_loss = reconstruction_loss
+        self.d_loss_tracker = tf.keras.metrics.Mean()
+        self.g_loss_tracker = tf.keras.metrics.Mean()
+        self.g_recon_tracker = tf.keras.metrics.Mean()
+        self.g_disc_tracker = tf.keras.metrics.Mean()
+
+        self.g_metric_trackers = [(tf.keras.metrics.Mean(), metric) for metric in reconstruction_metrics]
+        self.d_metric_trackers = [(tf.keras.metrics.Mean(), tf.keras.metrics.Mean(), tf.keras.metrics.Mean(), metric) for metric in discriminative_metrics]
+
+        all_trackers = [self.d_loss_tracker, self.g_loss_tracker, self.g_recon_tracker, self.g_disc_tracker] + \
+                       [tracker for tracker,_ in self.g_metric_trackers] + \
+                       [tracker for t in self.d_metric_trackers for tracker in t[:-1]]
+        self.all_trackers = MultipleTrackers(all_trackers)
+
+    def discriminator_loss(self, real_output, fake_output):
+        real_loss = self.loss(tf.ones_like(real_output), real_output)
+        fake_loss = self.loss(tf.zeros_like(fake_output), fake_output)
+        total_loss = 0.5*(real_loss + fake_loss)
+        return total_loss
+
+    def generator_loss(self, fake_output):
+        return self.loss(tf.ones_like(fake_output), fake_output)
+
+    @tf.function
+    def train_step(self, images):
+        masked, original = images
+        n_samples = tf.shape(original)[0]
+
+        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
+            generated_images = self.generator(masked, training=True)
+
+            real_output = self.discriminator(original, training=True)
+            fake_output = self.discriminator(generated_images, training=True)
+
+            gen_disc_loss = self.generator_loss(fake_output)
+            recon_loss = self.reconstruction_loss(original, generated_images)
+            gen_loss = self.C*recon_loss + gen_disc_loss
+            disc_loss = self.discriminator_loss(real_output, fake_output)
+
+        gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)
+        gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
+
+        self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
+        self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))
+
+        self.d_loss_tracker.update_state(tf.repeat([[disc_loss]], repeats=n_samples, axis=0))
+        self.g_loss_tracker.update_state(tf.repeat([[gen_loss]], repeats=n_samples, axis=0))
+        self.g_recon_tracker.update_state(tf.repeat([[recon_loss]], repeats=n_samples, axis=0))
+        self.g_disc_tracker.update_state(tf.repeat([[gen_disc_loss]], repeats=n_samples, axis=0))
+
+        logs = {'d_loss': self.d_loss_tracker.result()}
+
+        for tracker, real_tracker, fake_tracker, metric in self.d_metric_trackers:
+            v_real = metric(tf.ones_like(real_output), real_output)
+            v_fake = metric(tf.zeros_like(fake_output), fake_output)
+            v = 0.5*(v_real + v_fake)
+            tracker.update_state(tf.repeat([[v]], repeats=n_samples, axis=0))
+            real_tracker.update_state(tf.repeat([[v_real]], repeats=n_samples, axis=0))
+            fake_tracker.update_state(tf.repeat([[v_fake]], repeats=n_samples, axis=0))
+
+            metric_name = metric.__name__
+            logs['d_' + metric_name] = tracker.result()
+            logs['d_real_' + metric_name] = real_tracker.result()
+            logs['d_fake_' + metric_name] = fake_tracker.result()
+
+        logs['g_loss'] = self.g_loss_tracker.result()
+        logs['g_recon'] = self.g_recon_tracker.result()
+        logs['g_disc'] = self.g_disc_tracker.result()
+
+        for tracker, metric in self.g_metric_trackers:
+            v = metric(original, generated_images)
+            tracker.update_state(tf.repeat([[v]], repeats=n_samples, axis=0))
+            logs['g_' + metric.__name__] = tracker.result()
+
+        return logs
+
+    @tf.function
+    def val_step(self, images):
+        masked, original = images
+        n_samples = tf.shape(original)[0]
+
+        generated_images = self.generator(masked, training=False)
+
+        real_output = self.discriminator(original, training=False)
+        fake_output = self.discriminator(generated_images, training=False)
+
+        gen_disc_loss = self.generator_loss(fake_output)
+        recon_loss = self.reconstruction_loss(original, generated_images)
+        gen_loss = self.C*recon_loss + gen_disc_loss
+        disc_loss = self.discriminator_loss(real_output, fake_output)
+
+        self.d_loss_tracker.update_state(tf.repeat([[disc_loss]], repeats=n_samples, axis=0))
+        self.g_loss_tracker.update_state(tf.repeat([[gen_loss]], repeats=n_samples, axis=0))
+        self.g_recon_tracker.update_state(tf.repeat([[recon_loss]], repeats=n_samples, axis=0))
+        self.g_disc_tracker.update_state(tf.repeat([[gen_disc_loss]], repeats=n_samples, axis=0))
+
+        logs = {'val_d_loss': self.d_loss_tracker.result()}
+
+        for tracker, real_tracker, fake_tracker, metric in self.d_metric_trackers:
+            v_real = metric(tf.ones_like(real_output), real_output)
+            v_fake = metric(tf.zeros_like(fake_output), fake_output)
+            v = 0.5*(v_real + v_fake)
+            tracker.update_state(tf.repeat([[v]], repeats=n_samples, axis=0))
+            real_tracker.update_state(tf.repeat([[v_real]], repeats=n_samples, axis=0))
+            fake_tracker.update_state(tf.repeat([[v_fake]], repeats=n_samples, axis=0))
+
+            metric_name = metric.__name__
+            logs['val_d_' + metric_name] = tracker.result()
+            logs['val_d_real_' + metric_name] = real_tracker.result()
+            logs['val_d_fake_' + metric_name] = fake_tracker.result()
+
+        logs['val_g_loss'] = self.g_loss_tracker.result()
+        logs['val_g_recon'] = self.g_recon_tracker.result()
+        logs['val_g_disc'] = self.g_disc_tracker.result()
+
+        for tracker, metric in self.g_metric_trackers:
+            v = metric(original, generated_images)
+            tracker.update_state(tf.repeat([[v]], repeats=n_samples, axis=0))
+            logs['val_g_' + metric.__name__] = tracker.result()
+
+        return logs
+
+    def fit(self,
+            trainset,
+            valset=None,
+            trainsize=-1,
+            valsize=-1,
+            epochs=1,
+            display_per_epochs=5,
+            generator_callbacks=[],
+            discriminator_callbacks=[]):
+
+        print('🌊🐉 Start Training 🐉🌊')
+        gen_callback_tracker = tf.keras.callbacks.CallbackList(
+            generator_callbacks, add_history=True, model=self.generator
+        )
+
+        disc_callback_tracker = tf.keras.callbacks.CallbackList(
+            discriminator_callbacks, add_history=True, model=self.discriminator
+        )
+
+        callbacks_tracker = MultipleTrackers([gen_callback_tracker, disc_callback_tracker])
+
+        logs = {}
+        callbacks_tracker.on_train_begin(logs=logs)
+
+        for epoch in range(epochs):
+            print(f'Epochs {epoch+1}/{epochs}:')
+            callbacks_tracker.on_epoch_begin(epoch, logs=logs)
+
+            batches = tqdm(trainset,
+                           desc="Train",
+                           total=trainsize,
+                           unit="step",
+                           position=0,
+                           leave=True)
+
+            for batch, image_batch in enumerate(batches):
+
+                callbacks_tracker.on_batch_begin(batch, logs=logs)
+                callbacks_tracker.on_train_batch_begin(batch, logs=logs)
+
+                train_logs = {k:v.numpy() for k, v in self.train_step(image_batch).items()}
+                logs.update(train_logs)
+
+                callbacks_tracker.on_train_batch_end(batch, logs=logs)
+                callbacks_tracker.on_batch_end(batch, logs=logs)
+                batches.set_postfix({'d_loss': train_logs['d_loss'],
+                                     'g_loss': train_logs['g_loss']
+                                    })
+
+                # Presentation
+            stats = ", ".join("{}={:.3g}".format(k, v) for k, v in logs.items() if 'val_' not in k and 'loss' not in k)
+            print('Train:', stats)
+
+            batches.close()
+            if valset:
+                self.all_trackers.reset_state()
+
+                batches = tqdm(valset,
+                               desc="Valid",
+                               total=valsize,
+                               unit="step",
+                               position=0,
+                               leave=True)
+
+                for batch, image_batch in enumerate(batches):
+                    callbacks_tracker.on_batch_begin(batch, logs=logs)
+                    callbacks_tracker.on_test_batch_begin(batch, logs=logs)
+                    val_logs = {k:v.numpy() for k, v in self.val_step(image_batch).items()}
+                    logs.update(val_logs)
+
+                    callbacks_tracker.on_test_batch_end(batch, logs=logs)
+                    callbacks_tracker.on_batch_end(batch, logs=logs)
+                    # Presentation
+                    batches.set_postfix({'val_d_loss': val_logs['val_d_loss'],
+                                         'val_g_loss': val_logs['val_g_loss']
+                                        })
+
+                stats = ", ".join("{}={:.3g}".format(k, v) for k, v in logs.items() if 'val_' in k and 'loss' not in k)
+                print('Valid:', stats)
+
+                batches.close()
+
+            if epoch % display_per_epochs == 0:
+                print('-'*128)
+                self.visualize_samples((image_batch[0][:2], image_batch[1][:2]))
+
+            self.all_trackers.reset_state()
+
+            callbacks_tracker.on_epoch_end(epoch, logs=logs)
+#             tf.keras.backend.clear_session()
+            _ = gc.collect()
+
+            if self.generator.stop_training or self.discriminator.stop_training:
+                break
+            print('-'*128)
+
+        callbacks_tracker.on_train_end(logs=logs)
+        tf.keras.backend.clear_session()
+        _ = gc.collect()
+        gen_history = None
+        for cb in gen_callback_tracker:
+            if isinstance(cb, tf.keras.callbacks.History):
+                gen_history = cb
+                gen_history.history = {k:v for k,v in cb.history.items() if 'd_' not in k}
+
+        disc_history = None
+        for cb in disc_callback_tracker:
+            if isinstance(cb, tf.keras.callbacks.History):
+                disc_history = cb
+                disc_history.history = {k:v for k,v in cb.history.items() if 'g_' not in k}
+
+        return {'generator':gen_history,
+                'discriminator':disc_history}
+
+    def visualize_samples(self, samples, figsize=(12, 2)):
+        x, y = samples
+        y_pred = self.generator.predict(x[:2], verbose=0)
+        fig, axs = plt.subplots(1, 6, figsize=figsize)
+        for i in range(2):
+            pos = 3*i
+            axs[pos].imshow(x[i], cmap='gray', vmin=0., vmax=1.)
+            axs[pos].set_title('Masked')
+            axs[pos].axis('off')
+            axs[pos+1].imshow(y[i], cmap='gray', vmin=0., vmax=1.)
+            axs[pos+1].set_title('Original')
+            axs[pos+1].axis('off')
+            axs[pos+2].imshow(y_pred[i], cmap='gray', vmin=0., vmax=1.)
+            axs[pos+2].set_title('Predicted')
+            axs[pos+2].axis('off')
+        plt.show()
+
+#         tf.keras.backend.clear_session()
+        del y_pred
+        _ = gc.collect()

gae/modules.py

@@ -0,0 +1,241 @@
+import tensorflow as tf
+from tensorflow.keras.__internal__.layers import BaseRandomLayer
+from tensorflow.keras.layers import (
+    Dense, Flatten, Conv2D, Activation, BatchNormalization,
+    MaxPooling2D, AveragePooling2D, GlobalAveragePooling2D,
+    Dropout, Input, concatenate, add, Conv2DTranspose, Lambda,
+    SpatialDropout2D, Cropping2D, UpSampling2D, LeakyReLU,
+    ZeroPadding2D, Reshape, Concatenate, Multiply, Permute, Add
+)
+from keras import backend as K
+
+from .utils import normalize_tuple
+
+
+class MultipleTrackers():
+    def __init__(self, callback_lists: list):
+        self.callbacks_list = callback_lists
+
+    def __getattr__(self, attr):
+        def helper(*arg, **kwarg):
+            for cb in self.callbacks_list:
+                getattr(cb, attr)(*arg, **kwarg)
+        if attr in self.__class__.__dict__:
+            return getattr(self, attr)
+        else:
+            return helper
+
+
+class DropBlockNoise(BaseRandomLayer):
+    def __init__(
+        self,
+        rate,
+        block_size,
+        seed=None,
+        **kwargs,
+    ):
+        super().__init__(seed=seed, **kwargs)
+        if not 0.0 <= rate <= 1.0:
+            raise ValueError(
+                f"rate must be a number between 0 and 1. " f"Received: {rate}"
+            )
+
+        self._rate = rate
+        (
+            self._dropblock_height,
+            self._dropblock_width,
+        ) = normalize_tuple(
+            value=block_size, n=2, name="block_size", allow_zero=False
+        )
+        self.seed = seed
+
+    def call(self, x, training=None):
+        if not training or self._rate == 0.0:
+            return x
+
+        _, height, width, _ = tf.split(tf.shape(x), 4)
+
+        # Unnest scalar values
+        height = tf.squeeze(height)
+        width = tf.squeeze(width)
+
+        dropblock_height = tf.math.minimum(self._dropblock_height, height)
+        dropblock_width = tf.math.minimum(self._dropblock_width, width)
+
+        gamma = (
+            self._rate
+            * tf.cast(width * height, dtype=tf.float32)
+            / tf.cast(dropblock_height * dropblock_width, dtype=tf.float32)
+            / tf.cast(
+                (width - self._dropblock_width + 1)
+                * (height - self._dropblock_height + 1),
+                tf.float32,
+            )
+        )
+
+        # Forces the block to be inside the feature map.
+        w_i, h_i = tf.meshgrid(tf.range(width), tf.range(height))
+        valid_block = tf.logical_and(
+            tf.logical_and(
+                w_i >= int(dropblock_width // 2),
+                w_i < width - (dropblock_width - 1) // 2,
+            ),
+            tf.logical_and(
+                h_i >= int(dropblock_height // 2),
+                h_i < width - (dropblock_height - 1) // 2,
+            ),
+        )
+
+        valid_block = tf.reshape(valid_block, [1, height, width, 1])
+
+        random_noise = self._random_generator.random_uniform(
+            tf.shape(x), dtype=tf.float32
+        )
+        valid_block = tf.cast(valid_block, dtype=tf.float32)
+        seed_keep_rate = tf.cast(1 - gamma, dtype=tf.float32)
+        block_pattern = (1 - valid_block + seed_keep_rate + random_noise) >= 1
+        block_pattern = tf.cast(block_pattern, dtype=tf.float32)
+
+        window_size = [1, self._dropblock_height, self._dropblock_width, 1]
+
+        # Double negative and max_pool is essentially min_pooling
+        block_pattern = -tf.nn.max_pool(
+            -block_pattern,
+            ksize=window_size,
+            strides=[1, 1, 1, 1],
+            padding="SAME",
+        )
+
+        return (
+            x * tf.cast(block_pattern, x.dtype)
+        )
+
+
+def squeeze_excite_block(input, ratio=16):
+    ''' Create a channel-wise squeeze-excite block
+
+    Args:
+        input: input tensor
+        filters: number of output filters
+
+    Returns: a keras tensor
+
+    References
+    -   [Squeeze and Excitation Networks](https://arxiv.org/abs/1709.01507)
+    '''
+    init = input
+    channel_axis = 1 if K.image_data_format() == "channels_first" else -1
+    filters = int(init.shape[channel_axis])
+    se_shape = (1, 1, filters)
+
+    se = GlobalAveragePooling2D()(init)
+    se = Reshape(se_shape)(se)
+    se = Dense(filters // ratio, activation='relu', kernel_initializer='he_normal', use_bias=False)(se)
+    se = Dense(filters, activation='sigmoid', kernel_initializer='he_normal', use_bias=False)(se)
+
+    if K.image_data_format() == 'channels_first':
+        se = Permute((3, 1, 2))(se)
+
+    x = Multiply()([init, se])
+    return x
+
+
+def spatial_squeeze_excite_block(input):
+    ''' Create a spatial squeeze-excite block
+
+    Args:
+        input: input tensor
+
+    Returns: a keras tensor
+
+    References
+    -   [Concurrent Spatial and Channel Squeeze & Excitation in Fully Convolutional Networks](https://arxiv.org/abs/1803.02579)
+    '''
+
+    se = Conv2D(1, (1, 1), activation='sigmoid', use_bias=False,
+                kernel_initializer='he_normal')(input)
+
+    x = Multiply()([input, se])
+    return x
+
+
+def channel_spatial_squeeze_excite(input, ratio=16):
+    ''' Create a spatial squeeze-excite block
+
+    Args:
+        input: input tensor
+        filters: number of output filters
+
+    Returns: a keras tensor
+
+    References
+    -   [Squeeze and Excitation Networks](https://arxiv.org/abs/1709.01507)
+    -   [Concurrent Spatial and Channel Squeeze & Excitation in Fully Convolutional Networks](https://arxiv.org/abs/1803.02579)
+    '''
+
+    cse = squeeze_excite_block(input, ratio)
+    sse = spatial_squeeze_excite_block(input)
+
+    x = Add()([cse, sse])
+    return x
+
+
+def DoubleConv(filters, kernel_size, initializer='glorot_uniform'):
+    def layer(x):
+
+        x = Conv2D(filters, kernel_size, padding='same', kernel_initializer=initializer)(x)
+        x = BatchNormalization()(x)
+        x = Activation('swish')(x)
+        x = Conv2D(filters, kernel_size, padding='same', kernel_initializer=initializer)(x)
+        x = BatchNormalization()(x)
+        x = Activation('swish')(x)
+
+        return x
+
+    return layer
+
+
+def UpSampling2D_block(filters, kernel_size=(3, 3), upsample_rate=(2, 2), interpolation='bilinear',
+                       initializer='glorot_uniform', skip=None):
+    def layer(input_tensor):
+
+        x = UpSampling2D(size=upsample_rate, interpolation=interpolation)(input_tensor)
+
+        if skip is not None:
+            x = Concatenate()([x, skip])
+
+        x = DoubleConv(filters, kernel_size, initializer=initializer)(x)
+        x = channel_spatial_squeeze_excite(x)
+        return x
+
+    return layer
+
+
+def Conv2DTranspose_block(filters, transpose_kernel_size=(3, 3), upsample_rate=(2, 2),
+                          initializer='glorot_uniform', skip=None, met_input=None, sat_input=None):
+    def layer(input_tensor):
+        x = Conv2DTranspose(filters, transpose_kernel_size, strides=upsample_rate, padding='same')(input_tensor)
+        if skip is not None:
+            x = Concatenate()([x, skip])
+
+        x = DoubleConv(filters, transpose_kernel_size, initializer=initializer)(x)
+        x = channel_spatial_squeeze_excite(x)
+        return x
+
+    return layer
+
+
+def PixelShuffle_block(filters, kernel_size=(3, 3), upsample_rate=2,
+                          initializer='glorot_uniform', skip=None, met_input=None, sat_input=None):
+    def layer(input_tensor):
+        x = Conv2D(filters * (upsample_rate ** 2), kernel_size, padding="same",
+                   activation="swish", kernel_initializer='Orthogonal')(input_tensor)
+        x = tf.nn.depth_to_space(x, upsample_rate)
+        if skip is not None:
+            x = Concatenate()([x, skip])
+
+        x = DoubleConv(filters, kernel_size, initializer=initializer)(x)
+        x = channel_spatial_squeeze_excite(x)
+        return x
+
+    return layer

gae/utils.py

@@ -0,0 +1,180 @@
+import cv2
+import numpy as np
+from scipy import ndimage
+import matplotlib.pyplot as plt
+
+from skimage import exposure
+
+import tensorflow as tf
+import tensorflow_addons as tfa
+
+
+IMAGE_SIZE = 224
+
+
+def acc(y_true, y_pred, threshold=0.5):
+    threshold = tf.cast(threshold, y_pred.dtype)
+    y_pred = tf.cast(y_pred > threshold, y_pred.dtype)
+    return tf.reduce_mean(tf.cast(tf.equal(y_true, y_pred), tf.float32))
+
+def mae(y_true, y_pred):
+    return tf.reduce_mean(tf.abs(y_true-y_pred))
+
+def inv_ssim(y_true, y_pred):
+    return 1 - tf.reduce_mean(tf.image.ssim(y_true, y_pred, 1.0))
+
+def inv_msssim(y_true, y_pred):
+    return 1 - tf.reduce_mean(tf.image.ssim_multiscale(y_true, y_pred, 1.0, filter_size=4))
+
+def inv_msssim_l1(y_true, y_pred, alpha=0.8):
+    return alpha*inv_msssim(y_true, y_pred) + (1-alpha)*mae(y_true, y_pred)
+
+def inv_msssim_gaussian_l1(y_true, y_pred, alpha=0.8):
+    l1_diff = tf.abs(y_true-y_pred)
+    gaussian_l1 = tfa.image.gaussian_filter2d(l1_diff, filter_shape=(11, 11), sigma=1.5)
+    return alpha*inv_msssim(y_true, y_pred) + (1-alpha)*gaussian_l1
+
+def psnr(y_true, y_pred):
+    return tf.reduce_mean(tf.image.psnr(y_true, y_pred, 1.0))
+
+
+def show_image(image, title='Image', cmap_type='gray'):
+    plt.imshow(image, cmap=cmap_type)
+    plt.title(title)
+    plt.axis('off')
+    plt.show()
+
+
+# đảo màu những ảnh bị ngược màu
+def remove_negative(img):
+  outside = np.mean(img[ : , 0])
+  inside = np.mean(img[ : , int(IMAGE_SIZE / 2)])
+  if outside < inside:
+    return img
+  else:
+    return 1 - img
+
+# lựa chọn tiền xử lý: ảnh gốc, Equalization histogram, CLAHE
+def preprocess(img):
+    img = remove_negative(img)
+
+    img = exposure.equalize_hist(img)
+    img = exposure.equalize_adapthist(img)
+    img = exposure.equalize_hist(img)
+    return img
+
+
+# dilate contour
+def dilate(mask_img):
+    kernel_size = 2 * 22 + 1
+    kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
+    return ndimage.binary_dilation(mask_img == 0, structure=kernel)
+
+
+def normalize_tuple(value, n, name, allow_zero=False):
+    """Transforms non-negative/positive integer/integers into an integer tuple.
+    Args:
+      value: The value to validate and convert. Could an int, or any iterable of
+        ints.
+      n: The size of the tuple to be returned.
+      name: The name of the argument being validated, e.g. "strides" or
+        "kernel_size". This is only used to format error messages.
+      allow_zero: Default to False. A ValueError will raised if zero is received
+        and this param is False.
+    Returns:
+      A tuple of n integers.
+    Raises:
+      ValueError: If something else than an int/long or iterable thereof or a
+      negative value is
+        passed.
+    """
+    error_msg = (
+        f"The `{name}` argument must be a tuple of {n} "
+        f"integers. Received: {value}"
+    )
+
+    if isinstance(value, int):
+        value_tuple = (value,) * n
+    else:
+        try:
+            value_tuple = tuple(value)
+        except TypeError:
+            raise ValueError(error_msg)
+        if len(value_tuple) != n:
+            raise ValueError(error_msg)
+        for single_value in value_tuple:
+            try:
+                int(single_value)
+            except (ValueError, TypeError):
+                error_msg += (
+                    f"including element {single_value} of "
+                    f"type {type(single_value)}"
+                )
+                raise ValueError(error_msg)
+
+    if allow_zero:
+        unqualified_values = {v for v in value_tuple if v < 0}
+        req_msg = ">= 0"
+    else:
+        unqualified_values = {v for v in value_tuple if v <= 0}
+        req_msg = "> 0"
+
+    if unqualified_values:
+        error_msg += (
+            f" including {unqualified_values}"
+            f" that does not satisfy the requirement `{req_msg}`."
+        )
+        raise ValueError(error_msg)
+
+    return value_tuple
+
+
+def draw_points(image, points, color=None, random_color=False, same=True, thickness=1):
+    if color is None and not random_color:
+        color = (0, 255, 0)  # Màu mặc định là xanh lá cây (BGR)
+    if random_color:
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+    image = to_color(image)
+
+    for point in points:
+        if random_color and not same:
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+        x, y = point
+        image = cv2.circle(image, (x, y), thickness, color, -1)  # Vẽ điểm lên ảnh
+    return image
+
+
+def draw_lines(image, pairs, color=None, random_color=False, same=True, thickness=1):
+    image_with_line = to_color(np.copy(image))
+
+    if color is None and not random_color:
+        color = (0, 255, 0)  # Màu mặc định là xanh lá cây (BGR)
+    if random_color:
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+    # Vẽ đường thẳng dựa trên danh sách các cặp điểm
+    for pair in pairs:
+
+        if random_color and not same:
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+        start_point = pair[0]
+        end_point = pair[1]
+        image_with_line = cv2.line(image_with_line, start_point, end_point, color, thickness)
+        image_with_line = cv2.circle(image_with_line, start_point, thickness + 1, color, -1)
+        image_with_line = cv2.circle(image_with_line, end_point, thickness + 1, color, -1)
+
+    return image_with_line
+
+def to_color(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return image
+    return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+
+
+def to_gray(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    return image

gae/weights/gan_efficientunet_full_augment-hist_equal_generator.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6b3f57e8b8c2b6d56f1b0ac9eaceaa693600ccd52ba637293aee23eaf40a819
+size 230002208

jsw/__init__.py

@@ -0,0 +1 @@
+from .cal_jsw import get_JSW, calculate_diff, calculate_jsw_info

jsw/cal_jsw.py

@@ -0,0 +1,117 @@
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+from .utils import *
+from .contour import get_contours
+from .distance import distance
+
+
+def pooling_array(array, n, mode='mean'):
+    if mode == 'mean':
+        pool = lambda x: np.mean(x)
+    elif mode == 'min':
+        pool = lambda x: np.min(x)
+    elif mode == 'sum':
+        pool = lambda x: np.sum(x)
+
+    if n == 1:
+        return pool(array)
+
+    array_length = len(array)
+    if array_length < n:
+        return array
+    segment_length = array_length // n
+    remaining_elements = array_length % n
+
+    if remaining_elements == 0:
+        segments = np.split(array, n)
+    else:
+        mid = remaining_elements * (segment_length + 1)
+        segments = np.split(array[:mid], remaining_elements)
+        segments += np.split(array[mid:], n - remaining_elements)
+
+    segments = [pool(segment) for segment in segments]
+
+    return np.array(segments)
+
+
+def pool_links(links, dim, mode="mean"):
+    '''
+        links là 1 list gồm nhiều cặp tọa độ trên dưới (knee join space giữa lower_contour vs upper_contour) có dạng: 
+        [(array([436, 421], dtype=int32), array([436, 451], dtype=int32)), (array([436, 421], dtype=int32), array([436, 451], dtype=int32)), ...]
+        đầu tiên lấy danh sách các tọa độ x của các pair coord này, pool nó y như pool distance, rồi filter links ban đầu để lấy các pair có x thuộc pooled
+
+    '''
+    pooled_x = pooling_array(np.array(links)[:, 0, 0], dim, mode)
+    filtered_pairs = []
+    for x in pooled_x:
+        for pair in links:
+            if pair[0][0] == int(x):
+                filtered_pairs.append(pair)
+                break  # Lấy 1 cặp và dừng lại
+    return filtered_pairs
+
+
+def get_JSW(mask, dim=None, pool='mean', p=0.3, verbose=0):
+    '''
+        input: mask (H, W) with femur is 1 and tibia is 2
+        output: 2 distance vectors (left, right) (aggregated to <dim> by <pool>), links (list of pairs coords for each side)
+    '''
+    if isinstance(mask, str):
+        mask = cv2.imread(mask, 0)
+    if mask is None:
+        return np.zeros(10), np.zeros(10)
+    uc, lc = get_contours(mask, verbose=verbose)
+    left_distances, right_distances, links, contours = distance(mask, uc, lc, p=p, verbose=verbose)
+    if verbose:
+        # print('in getjsw')
+        temp = draw_points(mask * 127, contours[0], thickness=3, color=(255, 0, 0))
+        temp = draw_points(temp, contours[1], thickness=5, color=(255, 0, 0))
+        temp = draw_points(temp, contours[2], thickness=5, color=(0, 255, 0))
+        temp = draw_points(temp, contours[3], thickness=5, color=(0, 255, 0))
+        temp = draw_lines(temp, links[::6], color=(0, 0, 255), thickness=2)
+        # cv2_imshow(temp)
+        # cv2.imwrite("drawn_lines.png", temp)
+    if dim:
+        left_distances = pooling_array(left_distances, dim, pool)
+        right_distances = pooling_array(right_distances, dim, pool)
+        # print(left_distances)
+        # print(right_distances)
+
+        # print(links[1])
+        links = pool_links(links[0], dim, mode="mean"), pool_links(links[1], dim, mode="mean")
+    return left_distances, right_distances, links
+
+
+def calculate_diff(left_jsw, right_jsw):
+    '''
+        input: left_distances and right_distances vectors
+        output: jsw_m and jsw_mm
+    '''
+    jsw_max = max(np.max(left_jsw), np.max(right_jsw))
+    jsw_max_side = np.argmax([np.max(left_jsw), np.max(right_jsw)])
+    if jsw_max_side == 0:
+        jsw_min = np.min(right_jsw)
+    else:
+        jsw_min = np.min(left_jsw)
+    
+    diff_mean = abs(np.mean(left_jsw) - np.mean(right_jsw)) / jsw_max
+    diff_max_min = (jsw_max - jsw_min) / jsw_max
+
+    return diff_mean, diff_max_min
+
+
+def calculate_jsw_info(left_jsw, right_jsw):
+    '''
+        input: left_distances and right_distances vectors
+        output: % diff, mean_left, mean_right, side_min (0 is left, 1 is right), index with value min (in side min), value min (in side min)
+    '''
+    mean_left = np.mean(left_jsw)
+    mean_right = np.mean(right_jsw)
+
+    side_min, index_min, value_min = (0, np.argmin(left_jsw), np.min(left_jsw)) if mean_left <= mean_right else (1, np.argmin(right_jsw), np.min(right_jsw)) 
+
+    diff_percentage = np.abs((mean_left - mean_right) / mean_right) * 100
+
+    return diff_percentage, mean_left, mean_right, side_min, index_min, value_min

jsw/contour.py

@@ -0,0 +1,127 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+
+from .utils import *
+
+def detect_limit_points(mask, verbose=0):
+    # tìm giới hạn hai bên của khớp gối
+    '''
+        input: mask (H, W) with femur is 1 and tibia is 2
+        output: 4 points serve as limit points to determine the upper and lower contours from the full contours of the femur and tibia
+    '''
+    h, w = mask.shape
+    res = []
+    upper_pivot = np.array([0, w // 2])  # r c
+    lower_pivot = np.array([h, w // 2])  # r c
+
+    left_slice = slice(0, w // 2)
+    right_slice = slice(w // 2, None)
+    center_slice = slice(int(0.2 * h), int(0.8 * h))
+
+    left = np.zeros_like(mask)
+    left[center_slice, left_slice] = mask[center_slice, left_slice]
+
+    right = np.zeros_like(mask)
+    right[center_slice, right_slice] = mask[center_slice, right_slice]
+
+    if verbose:
+        cv2_imshow([left, right])
+
+    pivot = np.array([0, w])
+    coords = np.argwhere(left == 1)
+    distances = ((coords - pivot) ** 2).sum(axis=-1)
+    point = coords[distances.argmax()][::-1]
+    res.append(point)
+
+    pivot = np.array([0, 0])
+    coords = np.argwhere(right == 1)
+    distances = ((coords - pivot) ** 2).sum(axis=-1)
+    point = coords[distances.argmax()][::-1]
+    res.append(point)
+
+    pivot = np.array([h, w])
+    coords = np.argwhere(left == 2)
+    distances = ((coords - pivot) ** 2).sum(axis=-1)
+    point = coords[distances.argmax()][::-1]
+    res.append(point)
+
+    pivot = np.array([h, 0])
+    coords = np.argwhere(right == 2)
+    distances = ((coords - pivot) ** 2).sum(axis=-1)
+    point = coords[distances.argmax()][::-1]
+    res.append(point)
+
+    if verbose:
+        cv2_imshow(draw_points(127 * mask, res))
+
+    return res
+
+def find_boundaries(mask, start, end, top=True, verbose=0):
+    #     nếu top = True, tìm đường bao bên trên cùng từ left đến right
+    #     nếu top = False, tìm đường bao dưới cùng từ left đến right
+    '''
+        input:  
+            mask (H, W) of femur or tibia 
+            start, end is limit point to extract upper_contour/lower_contour from femur/tibia full contour
+        output: upper_contour/lower_contour
+        use top = True if determine lower contour from tibia mask
+    '''
+    boundaries = []
+    height, width = mask.shape
+
+    contours, _ = cv2.findContours(255 * mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+
+    areas = np.array([cv2.contourArea(cnt) for cnt in contours])
+    contour = contours[areas.argmax()]
+    contour = contour.reshape(-1, 2)
+    org_contour = contour.copy()
+
+    start_idx = ((start - contour) ** 2).sum(axis=-1).argmin()
+    end_idx = ((end - contour) ** 2).sum(axis=-1).argmin()
+    if start_idx <= end_idx:
+        contour = contour[start_idx:end_idx + 1]
+    else:
+        contour = np.concatenate([contour[start_idx:], contour[:end_idx + 1]])
+
+    if top:
+        sorted_indices = np.argsort(contour[:, 1])[::-1]
+    else:
+        sorted_indices = np.argsort(contour[:, 1])
+    contour = contour[sorted_indices]
+
+    unique_indices = sorted(np.unique(contour[:, 0], return_index=True)[1])
+    contour = contour[unique_indices]
+    sorted_indices = np.argsort(contour[:, 0])
+    contour = contour[sorted_indices]
+    if verbose:
+        temp = draw_points(127 * mask.astype(np.uint8), contour, thickness=5)
+        temp = draw_points(temp, [start, end], color=[155, 155], thickness=15)
+        # cv2_imshow(temp)
+
+    return np.array(contour), np.array(org_contour)
+
+
+def get_contours(mask, verbose=0):
+    '''
+        input: mask (H, W) with femur is 1 and tibia is 2
+        output: upper_contour, lower_contour
+    '''
+    limit_points = detect_limit_points(mask, verbose=verbose)
+    upper_contour, full_upper = find_boundaries(mask == 1, limit_points[0], limit_points[1], top=False, verbose=verbose)
+    lower_contour, full_lower = find_boundaries(mask == 2, limit_points[3], limit_points[2], top=True, verbose=verbose)
+    if verbose:
+        temp = draw_points(127 * mask, full_upper, thickness=3, color=(255, 0, 0))
+        temp = draw_points(temp, full_lower, thickness=3)
+        # cv2_imshow(temp)
+        # cv2.imwrite('full.png', temp)
+        temp = draw_points(temp, limit_points, thickness=7, color=(0, 0, 255))
+        # cv2_imshow(temp)
+        # cv2.imwrite('limit_points.png', temp)
+    if verbose:
+        temp = draw_points(127 * mask, upper_contour, thickness=3, color=(255, 0, 0))
+        temp = draw_points(temp, lower_contour, thickness=3)
+        # cv2_imshow(temp)
+        # cv2.imwrite('cropped.png', temp)
+
+    return upper_contour, lower_contour

jsw/distance.py

@@ -0,0 +1,98 @@
+from sklearn.linear_model import LinearRegression
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+from .utils import *
+
+
+def distance(mask, upper_contour, lower_contour, p=0.12, verbose = 0):
+    '''
+        input:
+            mask (H, W) with femur is 1 and tibia is 2
+            upper_contour, lower_contour
+        output: left_distances and right_distances vectors
+    '''
+    left, right = getMiddle(mask, lower_contour, verbose = verbose)
+    x_min = max(lower_contour[0,0],upper_contour[0,0])
+    x_max = min(lower_contour[-1,0],upper_contour[-1,0])
+    
+    left_idx = np.where(lower_contour[:,0] == left)[0][0]
+    right_idx = np.where(lower_contour[:,0] == right)[0][0]
+    left_lower_contour = lower_contour[left_idx:]
+    right_lower_contour = lower_contour[:right_idx+1][::-1]
+
+    left_lower_contour = lower_contour[(lower_contour[:,0]<=left) & (lower_contour[:,0]>=x_min) ]
+    right_lower_contour = lower_contour[(lower_contour[:,0]>=right) & (lower_contour[:,0]<=x_max)][::-1]
+
+    left_upper_contour = upper_contour[(upper_contour[:,0]<=left) & (upper_contour[:,0]>=x_min)]
+    right_upper_contour = upper_contour[(upper_contour[:,0]>=right) & (upper_contour[:,0]<=x_max)][::-1]
+
+    if verbose == 1:
+        temp = draw_points(mask*127, left_lower_contour, color = (0, 255, 0), thickness = 5)
+        temp = draw_points(temp, right_lower_contour, color = (0, 255, 0), thickness = 5)
+        temp = draw_points(temp, left_upper_contour, color = (255, 0, 0), thickness = 5)
+        temp = draw_points(temp, right_upper_contour, color = (255, 0, 0), thickness = 5)
+        # cv2_imshow(temp)
+        # cv2.imwrite('center_cropped.png', temp)
+    links =  list(zip(left_upper_contour,left_lower_contour)), list(zip(right_upper_contour,right_lower_contour))
+
+    temp = left_upper_contour, right_upper_contour, left_lower_contour, right_lower_contour
+
+    return left_lower_contour[:,1]-left_upper_contour[:,1], right_lower_contour[:,1] - right_upper_contour[:,1], links, temp
+
+
+def getMiddle(mask, contour, verbose=0):
+    '''
+        use Linear Regression to construct a straight line to remove the bone segment from upper_contour and lower_contour
+        input: 
+            mask (H, W) with femur is 1 and tibia is 2
+            lower_contour
+        output: 2 point (left and right) define the boundary within which every point of the bone segment is located
+    '''
+    X = contour[:, 0].reshape(-1, 1)
+    y = contour[:, 1]
+    reg = LinearRegression().fit(X, y)
+    i_min = np.argmin(y[int(len(y) * 0.2):int(len(y) * 0.8)]) + int(len(y) * 0.2)
+    left = i_min - 1
+    right = i_min + 1
+    left_check = False
+    right_check = False
+    if verbose == 1:
+        # print('get Middle 1')
+        cmask = draw_points(mask, contour, thickness=3, color=(255, 0, 0))
+        cmask = draw_points(cmask, np.hstack([X, reg.predict(X).reshape(-1, 1).astype('int')]),thickness = 5,
+                            color=(0, 255, 0))
+        # cv2.imwrite("lr_mask.png", cmask)
+        # cv2_imshow(cmask)
+        # plt.show()
+
+    while True:
+        while not left_check:
+            if y[left] > reg.predict(X[left].reshape(-1, 1)):
+                break
+            left -= 1
+        while not right_check:
+            if y[right] > reg.predict(X[right].reshape(-1, 1)):
+                break
+            right += 1
+        if verbose == 1:
+            # print('get middle 2')
+            cmask = draw_points(cmask, [contour[left]], thickness=10, color=(255, 255, 0))
+            cmask = draw_points(cmask, [contour[right]], thickness=10, color=(0, 255, 255))
+            # print(cmask.shape)
+            # cv2.imwrite("lr.png", cmask)
+            # cv2_imshow(cmask)
+            # plt.show()
+
+        left_min = np.argmin(y[int(len(y) * 0.2):left]) + int(len(y) * 0.2) if int(len(y) * 0.2) < left else left
+        right_min = np.argmin(y[right:int(len(y) * 0.8)]) + right if right < int(len(y) * 0.8) else right
+        if y[left_min] > reg.predict(X[left_min].reshape(-1, 1)):
+            left_check = True
+        if y[right_min] > reg.predict(X[right_min].reshape(-1, 1)):
+            right_check = True
+        if right_check and left_check:
+            break
+        left = left_min - 1
+        right = right_min + 1
+    return min(X.flatten()[left], X.flatten()[right]), max(X.flatten()[left], X.flatten()[right])

jsw/utils.py

@@ -0,0 +1,100 @@
+import random
+import cv2
+import matplotlib.pyplot as plt
+
+
+def draw_points(image, points, color=None, random_color=False, same=True, thickness=1):
+    if color is None and not random_color:
+        color = (0, 255, 0)  # Màu mặc định là xanh lá cây (BGR)
+    if random_color:
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+    image = to_color(image)
+
+    for point in points:
+        if random_color and not same:
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+        x, y = point
+        image = cv2.circle(image, (x, y), thickness, color, -1)  # Vẽ điểm lên ảnh
+    return image
+
+
+def draw_lines(image, pairs, color=None, random_color=False, same=True, thickness=1):
+    image_with_line = to_color(np.copy(image))
+
+    if color is None and not random_color:
+        color = (255, 0, 0)  # Màu mặc định là xanh lá cây (BGR)
+    if random_color:
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+    # Vẽ đường thẳng dựa trên danh sách các cặp điểm
+    for pair in pairs:
+
+        if random_color and not same:
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+        start_point = pair[0]
+        end_point = pair[1]
+        image_with_line = cv2.line(image_with_line, start_point, end_point, color, thickness)
+        image_with_line = cv2.circle(image_with_line, start_point, thickness + 1, color, -1)
+        image_with_line = cv2.circle(image_with_line, end_point, thickness + 1, color, -1)
+
+    return image_with_line
+
+
+def center(contour):
+    idx = len(contour) // 2
+    return contour[idx]
+
+
+def to_color(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return image
+    return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+
+
+def to_gray(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    return image
+
+
+def cv2_imshow(images):
+    if not isinstance(images, list):
+        images = [images]
+
+    num_images = len(images)
+
+    # Hiển thị ảnh đơn lẻ trực tiếp bằng imshow
+    if num_images == 1:
+        image = images[0]
+        if len(image.shape) == 3 and image.shape[2] == 3:
+            # Ảnh màu (RGB)
+            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            plt.imshow(image_rgb)
+        else:
+            # Ảnh xám
+            plt.imshow(image, cmap='gray')
+
+        plt.axis("off")
+        plt.show()
+    else:
+        # Hiển thị nhiều ảnh trên cùng một cột
+        fig, ax = plt.subplots(num_images, 1, figsize=(4, 4 * num_images))
+
+        for i in range(num_images):
+            image = images[i]
+            if len(image.shape) == 3 and image.shape[2] == 3:
+                # Ảnh màu (RGB)
+                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                ax[i].imshow(image_rgb)
+            else:
+                # Ảnh xám
+                ax[i].imshow(image, cmap='gray')
+
+            ax[i].axis("off")
+
+        plt.tight_layout()
+        plt.show()
+

requirements.txt

@@ -0,0 +1,19 @@
+streamlit==1.32.2
+opencv-python==4.9.0.80
+ultralytics
+numpy==1.26.4
+torch==2.2.2
+scikit-learn==1.2.2
+
+keras_cv~=0.5.0
+tqdm==4.66.2
+keras-core==0.1.7
+tensorflow==2.14
+tensorflow-addons==0.22.0
+scipy==1.12.0
+matplotlib==3.8.3
+scikit-image==0.22.0
+Pillow==10.2.0
+
+xgboost==2.1.2
+lime==0.2.0.1

segmentation/model.py

@@ -0,0 +1,37 @@
+from ultralytics import YOLO
+import cv2
+import torch
+
+BEST_WEIGHT = './segmentation/weights/oai_s_best4.pt'
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def create_model(weight_path=None):
+    if weight_path:
+        return YOLO(weight_path)
+    else:
+        return YOLO(BEST_WEIGHT)
+
+class Segmenter():
+    def __init__(self, weight_path=None):
+        self.model = create_model(weight_path).to(DEVICE)
+    
+    def segment(self, img):
+        """
+            input: image (H, W, C)
+            output: mask (H, W) with femur is 1 and tibia is 2 
+        """
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+        eimg = cv2.equalizeHist(img)
+
+        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+        eimg = clahe.apply(eimg)
+
+        eimg = cv2.cvtColor(eimg, cv2.COLOR_GRAY2RGB)
+
+        res = self.model(eimg, verbose=False)
+
+        mask = res[0].masks.data[0] * (res[0].boxes.cls[0] + 1) + res[0].masks.data[1] * (res[0].boxes.cls[1] + 1)
+        mask = mask.cpu().numpy()
+
+        return mask

segmentation/weights/oai_s_best4.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0159a91db899213c34db9cc93db1b5a31576eb3b78eb61bd1d9a29a4ef92e843
+size 6771320

utils.py

@@ -0,0 +1,79 @@
+import cv2
+import numpy as np
+# from shared_queue import image_queue
+
+
+def read_image(file_bytes):
+    np_img = np.frombuffer(file_bytes, np.uint8)
+    img = cv2.imdecode(np_img, cv2.IMREAD_COLOR)
+    return img
+
+def combine_mask(image, mask, label2color={1: (255, 255, 0), 2: (0, 255, 255)}, alpha=0.1):
+    image = to_color(image)
+    image = cv2.resize(image, mask.shape)
+    mask_image = np.zeros_like(image)
+    for label, color in label2color.items():
+        mask_image[mask == label] = color
+
+    mask_image = cv2.addWeighted(image, 1 - alpha, mask_image, alpha, 0)
+    return mask_image
+
+
+def to_color(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return image
+    return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+
+
+def to_gray(image):
+    if len(image.shape) == 3 and image.shape[-1] == 3:
+        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    return image
+
+def prob_process(prob_list):
+    prob = prob_list[0]
+    max_idx = np.argmax(prob)
+    min_idx = np.argmin(prob)
+    prob[max_idx] += prob[min_idx]
+    prob[min_idx] = 0
+    
+    if max_idx == 0:
+        return 1 - prob[max_idx]
+    elif max_idx ==1:
+        return prob[max_idx]
+    else:
+        return 1 + prob[max_idx]
+
+def combine_prob_jsw(prob, jsw_m, jsw_mm):
+    processed_prob = prob_process(prob)
+    temp = np.array([processed_prob, jsw_m, jsw_mm])
+    if temp.any():
+        temp = np.hstack(temp)
+    else:
+        temp = temp.flatten()
+    
+    return temp
+
+def scale_coordinates(links, original_size, mask_size):
+    scale_factor = original_size / mask_size
+    scaled_links = [(int(x * scale_factor), int(y * scale_factor)) for x, y in links]
+    return scaled_links
+
+def get_annotations(probabilites):
+    OSTEOPYTE_LEVELS = {
+        0: "Definite osteophytes",
+        1: "No osteophytes",
+        2: "Possible osteophytes",
+    }
+
+    JNS_LEVELS = {
+        0: "Definite JSN",
+        1: "Mild JSN",
+        2: "No JSN",
+        3: "Severe JSN",
+    }
+
+    return {
+        "osteophyte": OSTEOPYTE_LEVELS[np.argmax(probabilites[5:8])],
+        "jsn": JNS_LEVELS[np.argmax(probabilites[8:12])],
+    }