init

.idea/.gitignore

@@ -0,0 +1,8 @@
+# 默认忽略的文件
+/shelf/
+/workspace.xml
+# 基于编辑器的 HTTP 客户端请求
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/MOS2-train.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

.idea/deployment.xml

@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="PublishConfigData" remoteFilesAllowedToDisappearOnAutoupload="false">
+    <serverData>
+      <paths name="root@172.20.32.236:1522 password">
+        <serverdata>
+          <mappings>
+            <mapping local="$PROJECT_DIR$" web="/" />
+          </mappings>
+        </serverdata>
+      </paths>
+    </serverData>
+  </component>
+</project>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/MOS2-train.iml" filepath="$PROJECT_DIR$/.idea/MOS2-train.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

code/.ipynb_checkpoints/Visual-checkpoint.ipynb

@@ -0,0 +1,266 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "cb5a8f05-00bf-45b1-857d-6fb91ff72e4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from PIL import Image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "08a2e109-aaef-4ed7-b4ce-42c400afe3ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def dice(pred, label, num_classes=2, eps=1.0):\n",
+    "    \"\"\"\n",
+    "    计算多分类 Dice 系数，并返回平均值。\n",
+    "    pred: 预测的类别，大小为 (batch_size, height, width)。\n",
+    "    label: 真实标签类别，大小为 (batch_size, height, width)。\n",
+    "    num_classes: 类别数。\n",
+    "    \"\"\"\n",
+    "    dice_list = []\n",
+    "    \n",
+    "    for c in range(1, num_classes):\n",
+    "        # 对于每个类别，计算该类别的 Dice 系数\n",
+    "        pred_c = (pred == c)\n",
+    "        label_c = (label == c)\n",
+    "\n",
+    "        intersection = np.sum(pred_c & label_c)  # 交集\n",
+    "        union = np.sum(pred_c) + np.sum(label_c)  # 并集\n",
+    "\n",
+    "        if union / pred_c.size < 0.05:\n",
+    "            dice_list.append(1.)\n",
+    "            \n",
+    "        # if union == 0:\n",
+    "        #     dice_list.append(1.)\n",
+    "        \n",
+    "        # 计算 Dice 系数\n",
+    "        dice_list.append(2. * intersection / (union + eps))\n",
+    "\n",
+    "    return float(np.mean(dice_list))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "8ec7a6d9-8bf0-452c-af3c-4b65d9cfac1e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def softmax(x, axis=None):\n",
+    "    # 防止指数溢出（数值稳定技巧）\n",
+    "    x_max = np.max(x, axis=axis, keepdims=True)\n",
+    "    e_x = np.exp(x - x_max)\n",
+    "    return e_x / np.sum(e_x, axis=axis, keepdims=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "1eec453d-d503-4747-a626-b48bc24ecfd9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "exp_name = '0.25_unet'\n",
+    "json_path = './msunet/logs/{}/version_0/valid.json'.format(exp_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "f65061c6-1585-4d44-aeb0-9851f59d5559",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(json_path) as f:\n",
+    "    data = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cbb60a02-84bd-43fd-9cf5-b8ce98179ea1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "img_path = np.array(data['img_path'])\n",
+    "pred = np.array(data['pred']).argmax(1)\n",
+    "label = np.array(data['label'])\n",
+    "\n",
+    "num = len(img_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "73fdf7bd-91ec-4b9a-bd2b-7a6fffd632a8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "scores = [dice(pred[i], label[i]) for i in range(num)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "05760b40-2dc5-4715-bdc1-e0403156d7f0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.7910731675258925"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.mean(scores)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "018d5262-2333-465e-a7fa-2cdbfa883261",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "5f6c420a-126c-4424-babf-aaaf17be603b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.makedirs('../results/{}'.format(exp_name), exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "df284651-e024-4967-ad73-21996e747f76",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for i in range(num):\n",
+    "    img = Image.open(img_path[i])\n",
+    "    name = img_path[i].split('/')[-1].split('.')[0]\n",
+    "    plt.figure(figsize=(9, 3))\n",
+    "    plt.subplot(1, 3, 1)\n",
+    "    plt.imshow(img)\n",
+    "    plt.axis('off')\n",
+    "    plt.subplot(1, 3, 2)\n",
+    "    plt.imshow(label[i])\n",
+    "    plt.title('OTSU_Methods')\n",
+    "    plt.axis('off')\n",
+    "    plt.subplot(1, 3, 3)\n",
+    "    plt.imshow(pred[i])\n",
+    "    plt.title('UNet_{}'.format(int(round(scores[i], 2) * 100)))\n",
+    "    plt.axis('off')\n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig('../results/{}/{}.jpg'.format(exp_name, name))\n",
+    "    plt.close()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc7336f0-f003-49d7-9e27-0a4f2ab4caea",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2240a92f-fd25-4b1c-8faf-4474ed6a7c09",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dae57e32-8f23-43c2-9b57-a34ee0ea771e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "192f2672-99d1-46d1-a4e1-efc9bbae72bf",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58e3ec16-5874-4398-a1b2-64215042c0d9",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "850f7e49-6c4b-4f79-846e-558538c1a02c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90d749a4-ea51-4128-b88b-488cfd3318cd",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f8e6f1a5-b674-43c9-a70f-90d79814001f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/.ipynb_checkpoints/predict_pl-checkpoint.py

@@ -0,0 +1,98 @@
+import os
+import cv2
+import glob
+import torch
+import argparse
+import numpy as np
+import pandas as pd
+import albumentations as A
+import matplotlib.pyplot as plt
+
+from tqdm import tqdm
+from PIL import Image
+from albumentations.pytorch import ToTensorV2
+from torch.utils.data import Dataset, DataLoader
+
+from train_pl import SegModel
+from monai.inferers import SlidingWindowInferer
+
+
+# CONFIG
+def get_args():
+    parser = argparse.ArgumentParser(description="Inference configuration")
+
+    parser.add_argument("--dim",type=int,default=256,help="Feature dimension")
+    parser.add_argument("--model_output",type=str,default="./logs/",help="load pretrained model and save results")
+    parser.add_argument("--dataset", type=str, default="../data/test", help="Dataset path")
+
+    args = parser.parse_args()
+    return args
+
+
+class CustomDataset(Dataset):
+    def __init__(self, data_path):
+        self.imgs = glob.glob('{}/**/*.jpg'.format(data_path), recursive=True)
+        
+        self.aug = A.Compose([
+            A.Normalize(),
+            ToTensorV2(),
+        ])
+        
+    def __len__(self):
+        return len(self.imgs)
+    
+    def __getitem__(self, i):
+        image = np.array(Image.open(self.imgs[i]))
+        augmented = self.aug(image=image)
+        image = augmented['image'].float()
+        
+        return image
+
+def seg(args):
+    imgs = glob.glob('{}/**/*.jpg'.format(args.dataset), recursive=True)
+    ckpt_path = glob.glob('{}/**/*.ckpt'.format(args.model_output), recursive=True)[0]
+    model = SegModel.load_from_checkpoint(ckpt_path)
+    
+    model.eval();
+    
+    ds = CustomDataset(
+        data_path=args.dataset,
+    )
+    
+    dl = DataLoader(
+        dataset=ds,
+        batch_size=1,
+        num_workers=1,
+        shuffle=False,
+    )
+    
+    inferer = SlidingWindowInferer(
+        roi_size=(args.dim, args.dim),
+        sw_batch_size=4,
+        overlap=0.25,
+        padding_mode="reflect",
+        cache_roi_weight_map=True,
+        mode="gaussian"
+    )
+
+    save_path = '{}/test/'.format(args.model_output)
+    os.makedirs(save_path, exist_ok=True)
+    
+    for idx, image in enumerate(dl):
+        image = image.cuda()
+        
+        with torch.no_grad():
+            output = inferer(image, model)
+        
+        output = output.detach().cpu().numpy()
+        output = np.argmax(output, axis=1)
+        output[output == 1] = 255
+
+        output = np.array(output, np.uint8)
+        Image.fromarray(output[0]).save('{}/{}.png'.format(save_path, imgs[idx].split('/')[-1].split('.')[0]))
+            
+# main
+if __name__ == '__main__':
+    args = get_args()
+    seg(args)
+    

code/.ipynb_checkpoints/train_pl-checkpoint.py

@@ -0,0 +1,193 @@
+import os
+import json
+import glob
+import torch
+import argparse
+
+torch.set_float32_matmul_precision('high')
+
+import numpy as np
+from sklearn.utils import shuffle
+from torch import nn
+from torch.utils.data import DataLoader
+from torch.utils.data.sampler import *
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+import pytorch_lightning as pl
+
+from core.model import *
+from core.data import *
+from core.metrics import *
+
+
+# pytorch lightning module
+def get_args():
+    parser = argparse.ArgumentParser(description="Training configuration")
+
+    # Training params
+    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
+    parser.add_argument("--epochs", type=int, default=1000, help="Number of training epochs")
+    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
+    parser.add_argument("--weight_decay", type=float, default=1e-5, help="Weight decay")
+    parser.add_argument("--rf", type=float, default=0.9, help="Reduction factor / ratio factor")
+    parser.add_argument("--num_workers", type=int, default=4, help="Number of dataloader workers")
+
+    # Model params
+    parser.add_argument("--dim", type=int, default=256, help="Feature dimension")
+    parser.add_argument("--num_classes", type=int, default=3, help="Number of classes")
+    parser.add_argument("--model", type=str, default="unet_resnet34", help="Model name")
+    parser.add_argument(
+        "--pretrained_ckpt_path",
+        type=str,
+        default="../ckpt/resnet34-333f7ec4.pth",
+        help="Pretrained checkpoint path"
+    )
+
+    # Callback / save params
+    parser.add_argument("--save_top_k", type=int, default=1, help="Save top-k checkpoints")
+    parser.add_argument("--early_stop", type=int, default=2, help="Early stopping patience")
+    parser.add_argument("--every_n_epochs", type=int, default=1, help="Checkpoint save frequency")
+    parser.add_argument("--model_output", type=str, default="./logs/", help="Log save path")
+
+    # Data params
+    parser.add_argument("--dataset", type=str, default="../data/", help="Image dataset path")
+
+    args = parser.parse_args()
+    return args
+
+
+class SegModel(pl.LightningModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+        self.args = args
+        self.model = SMPModelFactory(
+            model=args.model, 
+            encoder_weights_path=args.pretrained_ckpt_path, 
+            classes=args.num_classes
+        ).get_model()
+        self.loss = CustomLoss()
+
+    def forward(self, x):
+        out = self.model(x)
+        return out
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer, factor=self.args.rf, mode='max', patience=2, min_lr=0)
+        return {
+            'optimizer': optimizer,
+            'lr_scheduler': scheduler,
+            'monitor': 'val_iou'
+        }
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        pd = self.model(x)
+        loss = self.loss(pd, y)
+        train_iou = iou(pd, y)
+        self.log('train_loss', loss)
+        self.log('train_iou', train_iou, on_epoch=True, prog_bar=True, logger=True)
+
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        pd = self.model(x)
+        loss = self.loss(pd, y)
+        val_iou = iou(pd, y)
+        self.log('val_loss', loss)
+        self.log('val_iou', val_iou, on_epoch=True, prog_bar=True, logger=True)
+
+    def predict_step(self, batch, batch_idx):
+        x, lbl = batch
+        pred = self.model(x)
+
+        return pred, lbl
+
+
+# main
+
+if __name__ == '__main__':
+    args = get_args()
+    train_img_list = np.array(glob.glob('{}/{}/*.jpg'.format(args.dataset, 'train'))).tolist()
+    valid_img_list = np.array(glob.glob('{}/{}/*.jpg'.format(args.dataset, 'valid'))).tolist()
+    
+    print('Train nums: {}, Valid nums: {}.'.format(len(train_img_list), len(valid_img_list)))
+
+    train_dataset = CustomDataset(train_img_list, dim=args.dim, data_type='train')
+    valid_dataset = CustomDataset(valid_img_list, dim=args.dim, data_type='valid')
+    
+    train_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        drop_last=True,
+    )
+
+    valid_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+    )
+
+    model = SegModel(args)
+
+    logger = TensorBoardLogger(
+        name=''.format(args.model),
+        save_dir=args.model_output,
+    )
+
+    checkpoint_callback = ModelCheckpoint(
+        every_n_epochs=args.every_n_epochs,
+        save_top_k=args.save_top_k,
+        monitor='val_iou',
+        mode='max',
+        filename='{epoch}-{val_loss:.4f}-{val_iou:.4f}'
+    )
+
+    earlystop_callback = EarlyStopping(
+        monitor="val_iou",
+        mode="max",
+        min_delta=0.00,
+        patience=args.early_stop,
+    )
+
+    # training
+    trainer = pl.Trainer(
+        accelerator='gpu',
+        devices=1,
+        max_epochs=args.epochs,
+        logger=logger,
+        callbacks=[checkpoint_callback, earlystop_callback]
+    )
+
+    trainer.fit(
+        model,
+        train_loader,
+        valid_loader
+    )
+
+    # inference
+    predictions = trainer.predict(
+        model=model,
+        dataloaders=valid_loader,
+        ckpt_path='best'
+    )
+
+    preds = torch.squeeze(torch.concat([item[0] for item in predictions])).numpy().tolist()
+    labels = torch.squeeze(torch.concat([item[1] for item in predictions])).numpy().tolist()
+
+    results = {
+        'img_path': valid_img_list,
+        'pred': preds,
+        'label': labels,
+    }
+
+    results_json = json.dumps(results)
+    with open(os.path.join(trainer.log_dir, 'valid.json'), 'w+') as f:
+        f.write(results_json)
+        

code/core/.ipynb_checkpoints/data-checkpoint.py

@@ -0,0 +1,85 @@
+import os
+import cv2
+import torch
+import numpy as np
+import pandas as pd
+import torch.nn as nn
+import albumentations as A
+import scipy.ndimage as ndimage
+import torchvision.transforms as T
+from albumentations.pytorch import ToTensorV2
+
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+
+# define heavy augmentations
+def get_training_augmentation(dim):
+    train_transform = [
+        A.Resize(dim, dim),
+        A.HorizontalFlip(),
+        A.VerticalFlip(),
+        A.GaussianBlur(),
+        A.RandomRotate90(),
+        A.RandomBrightnessContrast(),
+        A.ShiftScaleRotate(),
+        A.PadIfNeeded(min_height=dim, min_width=dim, always_apply=True),
+        A.Normalize(),
+        ToTensorV2(),
+    ]
+
+    return A.Compose(train_transform)
+
+def get_validation_augmentation(dim):
+    test_transform = [
+        A.Resize(dim, dim),
+        A.PadIfNeeded(min_height=dim, min_width=dim, always_apply=True, border_mode=0, value=0),
+        A.Normalize(),
+        ToTensorV2(),
+    ]
+
+    return A.Compose(test_transform)
+
+
+
+class CustomDataset(Dataset):
+    def __init__(self, img_list, dim=256, data_type='train'):
+        self.images = img_list
+        self.masks = [item.replace('.jpg', '.png') for item in self.images]
+        self.aug = get_training_augmentation(dim) if data_type == 'train' else get_validation_augmentation(dim)
+        
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, i):
+        image = cv2.imread(self.images[i])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        mask = cv2.imread(self.masks[i], cv2.IMREAD_GRAYSCALE)
+
+        augmented = self.aug(image=image, mask=mask)
+        image, mask = augmented['image'].float(), augmented['mask'].long()
+
+        return image, mask
+
+
+class CustomDatasetVideo(Dataset):
+    def __init__(self, img_list):
+        self.images = img_list
+        self.masks = [item.replace('.jpg', '.png') for item in self.images]
+        
+        self.aug = A.Compose([
+            A.Normalize(),
+            ToTensorV2(),
+        ])
+        
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, i):
+        image = cv2.imread(self.images[i])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        mask = cv2.imread(self.masks[i], cv2.IMREAD_GRAYSCALE)
+
+        augmented = self.aug(image=image, mask=mask)
+        image, mask = augmented['image'].float(), augmented['mask'].long()
+
+        return image, mask

code/core/.ipynb_checkpoints/e2e-checkpoint.py

@@ -0,0 +1,130 @@
+import os
+import glob
+import json
+import scipy
+import numpy as np
+import pandas as pd
+
+def center_coords_to_bbox(gt_coord):
+    box_rwidth, box_rheight = 10, 10
+    gt_bbox = (
+        gt_coord[0] - box_rwidth,
+        gt_coord[0] + box_rwidth + 1,
+        gt_coord[1] - box_rheight,
+        gt_coord[1] + box_rheight + 1
+    )
+    return gt_bbox
+
+
+def get_coord_to_bboxes(gt_coordinates_dict):
+    gt_bboxes_list = []
+    for gt_coord in gt_coordinates_dict:
+        gt_bbox = center_coords_to_bbox(gt_coord)
+        gt_bboxes_list.append(gt_bbox)
+    return gt_bboxes_list
+
+
+def bbox_iou(bb1, bb2):
+    assert bb1[0] <= bb1[1]
+    assert bb1[2] <= bb1[3]
+    assert bb2[0] <= bb2[1]
+    assert bb2[2] <= bb2[3]
+
+    # determine the coordinates of the intersection rectangle
+    x_left = max(bb1[0], bb2[0])
+    y_top = max(bb1[2], bb2[2])
+    x_right = min(bb1[1], bb2[1])
+    y_bottom = min(bb1[3], bb2[3])
+
+    if x_right < x_left or y_bottom < y_top:
+        return 0.0
+
+    # The intersection of two axis-aligned bounding boxes is always an
+    # axis-aligned bounding box
+    intersection_area = (x_right - x_left) * (y_bottom - y_top)
+
+    # compute the area of both AABBs
+    bb1_area = (bb1[1] - bb1[0]) * (bb1[3] - bb1[2])
+    bb2_area = (bb2[1] - bb2[0]) * (bb2[3] - bb2[2])
+
+    # compute the intersection over union by taking the intersection
+    # area and dividing it by the sum of prediction + ground-truth
+    # areas - the interesection area
+    iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
+    assert iou >= 0.0
+    assert iou <= 1.0
+    return iou
+
+
+def match_bboxes(iou_matrix, IOU_THRESH=0.5):
+    n_true, n_pred = iou_matrix.shape
+    MIN_IOU = 0.0
+    MAX_DIST = 1.0
+
+    if n_pred > n_true:
+        # there are more predictions than ground-truth - add dummy rows
+        diff = n_pred - n_true
+        iou_matrix = np.concatenate((iou_matrix,
+                                     np.full((diff, n_pred), MIN_IOU)),
+                                    axis=0)
+
+    if n_true > n_pred:
+        # more ground-truth than predictions - add dummy columns
+        diff = n_true - n_pred
+        iou_matrix = np.concatenate((iou_matrix,
+                                     np.full((n_true, diff), MIN_IOU)),
+                                    axis=1)
+
+    # call the Hungarian matching
+    idxs_true, idxs_pred = scipy.optimize.linear_sum_assignment(1 - iou_matrix)
+
+    if (not idxs_true.size) or (not idxs_pred.size):
+        ious = np.array([])
+    else:
+        ious = iou_matrix[idxs_true, idxs_pred]
+
+    # remove dummy assignments
+    sel_pred = idxs_pred < n_pred
+    idx_pred_actual = idxs_pred[sel_pred]
+    idx_gt_actual = idxs_true[sel_pred]
+    ious_actual = iou_matrix[idx_gt_actual, idx_pred_actual]
+    sel_valid = (ious_actual > IOU_THRESH)
+    label = sel_valid.astype(int)
+
+    return idx_gt_actual[sel_valid], idx_pred_actual[sel_valid], ious_actual[sel_valid], label
+
+
+def eval_matches(gt_bboxes, pd_bboxes, iou_threshold):
+        iou_matrix = np.zeros((len(gt_bboxes), len(pd_bboxes))).astype(np.float32)
+
+        for gt_idx, gt_bbox in enumerate(gt_bboxes):
+            for pd_idx, pd_bbox in enumerate(pd_bboxes):
+                iou = bbox_iou(gt_bbox, pd_bbox)
+                iou_matrix[gt_idx, pd_idx] = iou
+                
+        idxs_true, idxs_pred, ious, labels = match_bboxes(iou_matrix, IOU_THRESH=iou_threshold)
+        return idxs_true, idxs_pred, ious, labels
+    
+def eval_metrics(n_matches, n_gt, n_pred):
+        precision = n_matches / n_pred if n_pred > 0 else 0.0
+        if n_gt == 0:
+            raise RuntimeError("No ground truth atoms???")
+        recall = n_matches / n_gt
+        
+        return precision, recall
+    
+    
+def get_metrics(gt, pred, iou_threshold):
+    h, w = np.where(gt != 0)
+    gt_coords = list(zip(h.flatten(), w.flatten()))
+    gt_bboxes = get_coord_to_bboxes(gt_coords)
+
+    h, w = np.where(pred != 0)
+    pd_coords = list(zip(h.flatten(), w.flatten()))
+    pd_bboxes = get_coord_to_bboxes(pd_coords)
+
+    idxs_true, idxs_pred, ious, labels = eval_matches(gt_bboxes, pd_bboxes, iou_threshold)
+    precision, recall = eval_metrics(n_matches=len(idxs_pred), n_gt=len(gt_coords), n_pred=len(pd_bboxes))
+    f1_score = 2 * (precision * recall) / (precision + recall) if precision + recall > 0 else 0
+    
+    return precision, recall, f1_score

code/core/.ipynb_checkpoints/e2e_metrics-checkpoint.py

@@ -0,0 +1,130 @@
+import os
+import glob
+import json
+import scipy
+import numpy as np
+import pandas as pd
+
+def center_coords_to_bbox(gt_coord):
+    box_rwidth, box_rheight = 10, 10
+    gt_bbox = (
+        gt_coord[0] - box_rwidth,
+        gt_coord[0] + box_rwidth + 1,
+        gt_coord[1] - box_rheight,
+        gt_coord[1] + box_rheight + 1
+    )
+    return gt_bbox
+
+
+def get_coord_to_bboxes(gt_coordinates_dict):
+    gt_bboxes_list = []
+    for gt_coord in gt_coordinates_dict:
+        gt_bbox = center_coords_to_bbox(gt_coord)
+        gt_bboxes_list.append(gt_bbox)
+    return gt_bboxes_list
+
+
+def bbox_iou(bb1, bb2):
+    assert bb1[0] <= bb1[1]
+    assert bb1[2] <= bb1[3]
+    assert bb2[0] <= bb2[1]
+    assert bb2[2] <= bb2[3]
+
+    # determine the coordinates of the intersection rectangle
+    x_left = max(bb1[0], bb2[0])
+    y_top = max(bb1[2], bb2[2])
+    x_right = min(bb1[1], bb2[1])
+    y_bottom = min(bb1[3], bb2[3])
+
+    if x_right < x_left or y_bottom < y_top:
+        return 0.0
+
+    # The intersection of two axis-aligned bounding boxes is always an
+    # axis-aligned bounding box
+    intersection_area = (x_right - x_left) * (y_bottom - y_top)
+
+    # compute the area of both AABBs
+    bb1_area = (bb1[1] - bb1[0]) * (bb1[3] - bb1[2])
+    bb2_area = (bb2[1] - bb2[0]) * (bb2[3] - bb2[2])
+
+    # compute the intersection over union by taking the intersection
+    # area and dividing it by the sum of prediction + ground-truth
+    # areas - the interesection area
+    iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
+    assert iou >= 0.0
+    assert iou <= 1.0
+    return iou
+
+
+def match_bboxes(iou_matrix, IOU_THRESH=0.5):
+    n_true, n_pred = iou_matrix.shape
+    MIN_IOU = 0.0
+    MAX_DIST = 1.0
+
+    if n_pred > n_true:
+        # there are more predictions than ground-truth - add dummy rows
+        diff = n_pred - n_true
+        iou_matrix = np.concatenate((iou_matrix,
+                                     np.full((diff, n_pred), MIN_IOU)),
+                                    axis=0)
+
+    if n_true > n_pred:
+        # more ground-truth than predictions - add dummy columns
+        diff = n_true - n_pred
+        iou_matrix = np.concatenate((iou_matrix,
+                                     np.full((n_true, diff), MIN_IOU)),
+                                    axis=1)
+
+    # call the Hungarian matching
+    idxs_true, idxs_pred = scipy.optimize.linear_sum_assignment(1 - iou_matrix)
+
+    if (not idxs_true.size) or (not idxs_pred.size):
+        ious = np.array([])
+    else:
+        ious = iou_matrix[idxs_true, idxs_pred]
+
+    # remove dummy assignments
+    sel_pred = idxs_pred < n_pred
+    idx_pred_actual = idxs_pred[sel_pred]
+    idx_gt_actual = idxs_true[sel_pred]
+    ious_actual = iou_matrix[idx_gt_actual, idx_pred_actual]
+    sel_valid = (ious_actual > IOU_THRESH)
+    label = sel_valid.astype(int)
+
+    return idx_gt_actual[sel_valid], idx_pred_actual[sel_valid], ious_actual[sel_valid], label
+
+
+def eval_matches(gt_bboxes, pd_bboxes, iou_threshold):
+        iou_matrix = np.zeros((len(gt_bboxes), len(pd_bboxes))).astype(np.float32)
+
+        for gt_idx, gt_bbox in enumerate(gt_bboxes):
+            for pd_idx, pd_bbox in enumerate(pd_bboxes):
+                iou = bbox_iou(gt_bbox, pd_bbox)
+                iou_matrix[gt_idx, pd_idx] = iou
+                
+        idxs_true, idxs_pred, ious, labels = match_bboxes(iou_matrix, IOU_THRESH=iou_threshold)
+        return idxs_true, idxs_pred, ious, labels
+    
+def eval_metrics(n_matches, n_gt, n_pred):
+        precision = n_matches / n_pred if n_pred > 0 else 0.0
+        if n_gt == 0:
+            raise RuntimeError("No ground truth atoms???")
+        recall = n_matches / n_gt
+        
+        return precision, recall
+    
+    
+def get_metrics(gt, pred, iou_threshold):
+    h, w = np.where(gt != 0)
+    gt_coords = list(zip(h.flatten(), w.flatten()))
+    gt_bboxes = get_coord_to_bboxes(gt_coords)
+
+    h, w = np.where(pred != 0)
+    pd_coords = list(zip(h.flatten(), w.flatten()))
+    pd_bboxes = get_coord_to_bboxes(pd_coords)
+
+    idxs_true, idxs_pred, ious, labels = eval_matches(gt_bboxes, pd_bboxes, iou_threshold)
+    precision, recall = eval_metrics(n_matches=len(idxs_pred), n_gt=len(gt_coords), n_pred=len(pd_bboxes))
+    f1_score = 2 * (precision * recall) / (precision + recall) if precision + recall > 0 else 0
+    
+    return precision, recall, f1_score

code/core/.ipynb_checkpoints/metrics-checkpoint.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+
+
+class CustomLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dice_loss = smp.losses.DiceLoss(mode="multiclass")
+        self.ce_loss = nn.CrossEntropyLoss()
+                
+    def forward(self, y_pred, y_true):
+        return self.dice_loss(y_pred, y_true) + self.ce_loss(y_pred, y_true)    
+
+
+# def dice(pred, target, eps=1.0):
+#     """
+#     计算每个类别的 Dice 系数，并返回平均值。
+#     pred: 预测的类别（预测的最大类别索引）。
+#     target: 真实标签类别（每个像素的类别索引）。
+#     """
+#     num_classes = pred.shape[1]
+#     pred = torch.argmax(pred, dim=1)
+    
+#     dice_list = []
+#     for c in range(1, num_classes):
+#         # 对于每个类别，计算该类别的 Dice 系数
+#         pred_c = (pred == c)
+#         target_c = (target == c)
+        
+#         intersection = torch.sum(pred_c & target_c)
+        
+#         # 计算 Dice 系数
+#         dice_list.append((2 * intersection) / (torch.sum(pred_c) + torch.sum(target_c) + eps))
+    
+#     return torch.mean(torch.tensor(dice_list))  # 返回所有类别的平均 Dice 系数
+
+def iou(pred, target, eps=1.0):
+    """
+    计算每个类别的 IoU 系数，并返回平均值。
+    pred: 预测的类别（预测的最大类别索引）。
+    target: 真实标签类别（每个像素的类别索引）。
+    """
+    num_classes = pred.shape[1]
+    pred = torch.argmax(pred, dim=1)
+    
+    iou_list = []
+    for c in range(1, num_classes):
+        # 对于每个类别，计算该类别的 IoU
+        pred_c = (pred == c)
+        target_c = (target == c)
+        
+        intersection = torch.sum(pred_c & target_c)
+        union = torch.sum(pred_c | target_c)
+
+        # 计算 IoU
+        if union != 0:
+            iou_list.append(intersection / (union + eps))
+    
+    return torch.mean(torch.tensor(iou_list))  # 返回所有类别的平均 IoU

code/core/.ipynb_checkpoints/model-checkpoint.py

@@ -0,0 +1,59 @@
+import torch
+import segmentation_models_pytorch as smp
+
+class SMPModelFactory:
+    def __init__(self, model="unet_resnet34", encoder_weights_path="", in_channels=3, classes=2):
+        self.model_name = model.split('_')[0].lower()
+        self.encoder_name = model.split('_')[1].lower()
+        self.encoder_weights_path = encoder_weights_path
+        self.in_channels = in_channels
+        self.classes = classes
+
+    def get_model(self):
+        if self.model_name == "unet":
+            model = smp.Unet(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "unet++":
+            model = smp.UnetPlusPlus(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "fpn":
+            model = smp.FPN(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "deeplabv3plus":
+            model = smp.DeepLabV3Plus(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "pspnet":
+            model = smp.PSPNet(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "linknet":
+            model = smp.Linknet(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        else:
+            raise ValueError(f"Unsupported model name: {model_name}")
+
+        model.encoder.load_state_dict(torch.load(self.encoder_weights_path, weights_only=False))
+        return model

code/core/data.py

@@ -0,0 +1,86 @@
+import os
+import cv2
+import torch
+import numpy as np
+import pandas as pd
+import torch.nn as nn
+import albumentations as A
+import scipy.ndimage as ndimage
+import torchvision.transforms as T
+from albumentations.pytorch import ToTensorV2
+
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+
+# define heavy augmentations
+def get_training_augmentation(dim):
+    train_transform = [
+        A.Resize(dim, dim),
+        A.HorizontalFlip(),
+        A.VerticalFlip(),
+        A.GaussianBlur(),
+        A.RandomRotate90(),
+        A.RandomBrightnessContrast(),
+        # A.ShiftScaleRotate(border_mode=cv2.BORDER_REFLECT_101),
+        A.ShiftScaleRotate(),
+        A.PadIfNeeded(min_height=dim, min_width=dim, always_apply=True),
+        A.Normalize(),
+        ToTensorV2(),
+    ]
+
+    return A.Compose(train_transform)
+
+def get_validation_augmentation(dim):
+    test_transform = [
+        A.Resize(dim, dim),
+        A.PadIfNeeded(min_height=dim, min_width=dim, always_apply=True, border_mode=0, value=0),
+        A.Normalize(),
+        ToTensorV2(),
+    ]
+
+    return A.Compose(test_transform)
+
+
+
+class CustomDataset(Dataset):
+    def __init__(self, img_list, dim=256, data_type='train'):
+        self.images = img_list
+        self.masks = [item.replace('.jpg', '.png') for item in self.images]
+        self.aug = get_training_augmentation(dim) if data_type == 'train' else get_validation_augmentation(dim)
+        
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, i):
+        image = cv2.imread(self.images[i])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        mask = cv2.imread(self.masks[i], cv2.IMREAD_GRAYSCALE)
+
+        augmented = self.aug(image=image, mask=mask)
+        image, mask = augmented['image'].float(), augmented['mask'].long()
+
+        return image, mask
+
+
+class CustomDatasetVideo(Dataset):
+    def __init__(self, img_list):
+        self.images = img_list
+        self.masks = [item.replace('.jpg', '.png') for item in self.images]
+        
+        self.aug = A.Compose([
+            A.Normalize(),
+            ToTensorV2(),
+        ])
+        
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, i):
+        image = cv2.imread(self.images[i])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        mask = cv2.imread(self.masks[i], cv2.IMREAD_GRAYSCALE)
+
+        augmented = self.aug(image=image, mask=mask)
+        image, mask = augmented['image'].float(), augmented['mask'].long()
+
+        return image, mask

code/core/metrics.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+
+
+class CustomLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dice_loss = smp.losses.DiceLoss(mode="multiclass")
+        self.ce_loss = nn.CrossEntropyLoss()
+                
+    def forward(self, y_pred, y_true):
+        return self.dice_loss(y_pred, y_true) + self.ce_loss(y_pred, y_true)    
+
+
+# def dice(pred, target, eps=1.0):
+#     """
+#     计算每个类别的 Dice 系数，并返回平均值。
+#     pred: 预测的类别（预测的最大类别索引）。
+#     target: 真实标签类别（每个像素的类别索引）。
+#     """
+#     num_classes = pred.shape[1]
+#     pred = torch.argmax(pred, dim=1)
+    
+#     dice_list = []
+#     for c in range(1, num_classes):
+#         # 对于每个类别，计算该类别的 Dice 系数
+#         pred_c = (pred == c)
+#         target_c = (target == c)
+        
+#         intersection = torch.sum(pred_c & target_c)
+        
+#         # 计算 Dice 系数
+#         dice_list.append((2 * intersection) / (torch.sum(pred_c) + torch.sum(target_c) + eps))
+    
+#     return torch.mean(torch.tensor(dice_list))  # 返回所有类别的平均 Dice 系数
+
+def iou(pred, target, eps=1.0):
+    """
+    计算每个类别的 IoU 系数，并返回平均值。
+    pred: 预测的类别（预测的最大类别索引）。
+    target: 真实标签类别（每个像素的类别索引）。
+    """
+    num_classes = pred.shape[1]
+    pred = torch.argmax(pred, dim=1)
+    
+    iou_list = []
+    for c in range(1, num_classes):
+        # 对于每个类别，计算该类别的 IoU
+        pred_c = (pred == c)
+        target_c = (target == c)
+        
+        intersection = torch.sum(pred_c & target_c)
+        union = torch.sum(pred_c | target_c)
+
+        # 计算 IoU
+        if union != 0:
+            iou_list.append(intersection / (union + eps))
+    
+    return torch.mean(torch.tensor(iou_list))  # 返回所有类别的平均 IoU

code/core/model.py

@@ -0,0 +1,59 @@
+import torch
+import segmentation_models_pytorch as smp
+
+class SMPModelFactory:
+    def __init__(self, model="unet_resnet34", encoder_weights_path="", in_channels=3, classes=2):
+        self.model_name = model.split('_')[0].lower()
+        self.encoder_name = model.split('_')[1].lower()
+        self.encoder_weights_path = encoder_weights_path
+        self.in_channels = in_channels
+        self.classes = classes
+
+    def get_model(self):
+        if self.model_name == "unet":
+            model = smp.Unet(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "unet++":
+            model = smp.UnetPlusPlus(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "fpn":
+            model = smp.FPN(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "deeplabv3plus":
+            model = smp.DeepLabV3Plus(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "pspnet":
+            model = smp.PSPNet(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        elif self.model_name == "linknet":
+            model = smp.Linknet(
+                encoder_name=self.encoder_name,
+                encoder_weights=None,
+                in_channels=self.in_channels,
+                classes=self.classes
+            )
+        else:
+            raise ValueError(f"Unsupported model name: {model_name}")
+
+        model.encoder.load_state_dict(torch.load(self.encoder_weights_path, weights_only=False))
+        return model