{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":113558,"databundleVersionId":14878066,"sourceType":"competition"},{"sourceId":14137814,"sourceType":"datasetVersion","datasetId":9009300},{"sourceId":14141160,"sourceType":"datasetVersion","datasetId":9011842},{"sourceId":14142451,"sourceType":"datasetVersion","datasetId":9012780},{"sourceId":285918305,"sourceType":"kernelVersion"}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Segformer","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\n!pip install segmentation-models-pytorch\n\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\n\ntry:\n    import google.protobuf.message_factory\n    from google.protobuf.message_factory import MessageFactory\n    \n    if not hasattr(MessageFactory, 'GetPrototype'):\n        print(\"Applying patch ..\")\n        \n        def GetPrototype(self, descriptor):\n            \n            if hasattr(self, 'pool'):\n                return self.pool.GetPrototype(descriptor)\n            return None\n            \n        MessageFactory.GetPrototype = GetPrototype\n        print(\"Patch applied\")\nexcept ImportError:\n    print(\"Error.\")\nexcept Exception as e:\n    print(f\"Error: {e}\")\n\n\nimport glob\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nimport json\nimport gc\nfrom torch import nn\nfrom tqdm.auto import tqdm\nfrom transformers import SegformerForSemanticSegmentation, SegformerImageProcessor\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\n\n\nclass ForgeryDataset(Dataset):\n    def __init__(self, image_paths, mask_dir, processor):\n        self.image_paths = image_paths\n        self.mask_dir = mask_dir\n        self.processor = processor\n        \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        file_name = os.path.basename(img_path)\n        \n        image = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)\n        if image is None: return self._get_dummy()\n\n        if image.ndim == 2: image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)\n        elif image.shape[2] == 4: image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)\n        else: image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n            \n        orig_h, orig_w = image.shape[:2]\n        \n        mask = np.zeros((orig_h, orig_w), dtype=np.float32)\n\n        if 'authentic' not in img_path:\n            mask_name = file_name.replace('.png', '.npy')\n            mask_path = os.path.join(self.mask_dir, mask_name)\n            \n            if os.path.exists(mask_path):\n                try:\n                    loaded_mask = np.load(mask_path)\n                    if loaded_mask.ndim == 3:\n                        if loaded_mask.shape[0] < 5: loaded_mask = np.max(loaded_mask, axis=0)\n                        elif loaded_mask.shape[-1] < 5: loaded_mask = np.max(loaded_mask, axis=-1)\n                    loaded_mask = np.squeeze(loaded_mask)\n                    if loaded_mask.shape != (orig_h, orig_w):\n                        loaded_mask = cv2.resize(loaded_mask.astype(np.float32), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)\n                    mask = loaded_mask.astype(np.float32)\n                except: pass\n\n        try:\n            inputs = self.processor(images=image, return_tensors=\"pt\")\n            pixel_values = inputs.pixel_values.squeeze()\n            \n            mask_resized = cv2.resize(mask, (Config.IMAGE_SIZE, Config.IMAGE_SIZE), interpolation=cv2.INTER_NEAREST)\n            if mask_resized.ndim > 2: mask_resized = mask_resized[:, :, 0]\n            mask_tensor = torch.tensor(mask_resized).float().unsqueeze(0)\n            \n            if pixel_values.shape != (3, Config.IMAGE_SIZE, Config.IMAGE_SIZE) or \\\n               mask_tensor.shape != (1, Config.IMAGE_SIZE, Config.IMAGE_SIZE):\n                   return self._get_dummy()\n\n            return {\n                \"pixel_values\": pixel_values,\n                \"labels\": mask_tensor,\n                \"original_size\": torch.tensor([orig_h, orig_w])\n            }\n        except: return self._get_dummy()\n\n    def _get_dummy(self):\n        return {\n            \"pixel_values\": torch.zeros((3, Config.IMAGE_SIZE, Config.IMAGE_SIZE)),\n            \"labels\": torch.zeros((1, Config.IMAGE_SIZE, Config.IMAGE_SIZE)),\n            \"original_size\": torch.tensor([512, 512])\n        }\n\nBAD_FILES_LIST = ['57085.png', '25311.png', '34265.png', '64128.png']\n\n\ndef run_training():\n    print(\"Process data...\")\n    authentic_imgs = glob.glob(os.path.join(Config.TRAIN_IMG_AUTH, \"*.png\"))\n    forged_imgs = glob.glob(os.path.join(Config.TRAIN_IMG_FORGED, \"*.png\"))\n    forged_imgs = [x for x in forged_imgs if os.path.basename(x) not in BAD_FILES_LIST]\n    \n    all_train_images = authentic_imgs + forged_imgs\n    np.random.shuffle(all_train_images)\n    \n    split_idx = int(len(all_train_images) * 0.9)\n    train_files = all_train_images[:split_idx]\n    val_files = all_train_images[split_idx:]\n    \n    processor = SegformerImageProcessor.from_pretrained(\n        Config.MODEL_PATH, \n        do_resize=True, \n        size={\"height\": Config.IMAGE_SIZE, \"width\": Config.IMAGE_SIZE}\n    )\n    \n    train_dataset = ForgeryDataset(train_files, Config.TRAIN_MASKS, processor)\n    val_dataset = ForgeryDataset(val_files, Config.TRAIN_MASKS, processor)\n    \n    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2)\n    \n    print(f\"Load models: {Config.MODEL_PATH}\")\n    model = SegformerForSemanticSegmentation.from_pretrained(\n        Config.MODEL_PATH, \n        num_labels=1, \n        ignore_mismatched_sizes=True\n    )\n    model.to(Config.DEVICE)\n    \n    optimizer = AdamW(model.parameters(), lr=Config.LR)\n    criterion = nn.BCEWithLogitsLoss()\n    \n    print(\"Start Training...\")\n    for epoch in range(Config.EPOCHS):\n        model.train()\n        train_loss = 0\n        progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{Config.EPOCHS}\")\n        \n        for batch in progress_bar:\n            pixel_values = batch[\"pixel_values\"].to(Config.DEVICE)\n            labels = batch[\"labels\"].to(Config.DEVICE)\n            \n            optimizer.zero_grad()\n            outputs = model(pixel_values=pixel_values)\n            \n            logits_upsampled = nn.functional.interpolate(\n                outputs.logits, size=(Config.IMAGE_SIZE, Config.IMAGE_SIZE), \n                mode=\"bilinear\", align_corners=False\n            )\n            \n            loss = criterion(logits_upsampled, labels)\n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item()\n            progress_bar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})\n            \n        print(f\"Epoch {epoch+1} done. Loss: {train_loss/len(train_loader):.4f}\")\n        \n    print(f\"Save model into {Config.SAVE_DIR}...\")\n    model.save_pretrained(Config.SAVE_DIR)\n    processor.save_pretrained(Config.SAVE_DIR)\n    print(\"Model saved\")\n\n\n\nclass Config:\n    ROOT_DIR = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n    TEST_IMGS = os.path.join(ROOT_DIR, \"test_images\")\n    \n    TRAIN_MODE = False \n    \n    MODEL_PATH = \"/kaggle/input/my-segformer-weights/saved_model\" \n    \n    IMAGE_SIZE = 512\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\ndef rle_encode(mask):\n    pixels = mask.flatten(order='F')\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return json.dumps(runs.tolist())\n\n\ndef run_inference():\n    print(f\"Load model from: {Config.MODEL_PATH}\")\n    \n    if not os.path.exists(Config.MODEL_PATH):\n        print(f\"Error path {Config.MODEL_PATH} do not exist!\")\n        return\n\n    try:\n        processor = SegformerImageProcessor.from_pretrained(Config.MODEL_PATH)\n        model = SegformerForSemanticSegmentation.from_pretrained(Config.MODEL_PATH)\n    except Exception as e:\n        print(f\"Error loading model {e}\")\n        return\n\n    model.to(Config.DEVICE)\n    model.eval()\n    \n    submission_data = []\n    \n    test_extensions = [\"*.png\", \"*.jpg\", \"*.jpeg\", \"*.tif\", \"*.tiff\"]\n    test_images = []\n    for ext in test_extensions:\n        test_images.extend(glob.glob(os.path.join(Config.TEST_IMGS, ext)))\n        test_images.extend(glob.glob(os.path.join(Config.TEST_IMGS, ext.upper())))\n    \n    test_images = sorted(list(set(test_images)))\n    print(f\"Found imgages: {len(test_images)}\")\n    \n    with torch.no_grad():\n        for img_path in tqdm(test_images):\n            filename = os.path.basename(img_path)\n            case_id = os.path.splitext(filename)[0]\n            \n            try:\n                image = cv2.imread(img_path)\n                if image is None: raise ValueError(\"Read error\")\n                \n                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n                orig_h, orig_w = image.shape[:2]\n                \n                inputs = processor(images=image, return_tensors=\"pt\")\n                pixel_values = inputs.pixel_values.to(Config.DEVICE)\n                \n                outputs = model(pixel_values=pixel_values)\n                \n                logits_upsampled = nn.functional.interpolate(\n                    outputs.logits, size=(orig_h, orig_w), \n                    mode=\"bilinear\", align_corners=False\n                )\n                \n                probs = torch.sigmoid(logits_upsampled).cpu().numpy()[0, 0]\n                pred_mask = (probs > 0.5).astype(np.uint8)\n                \n                if np.sum(pred_mask) == 0:\n                    annotation = \"authentic\"\n                else:\n                    annotation = rle_encode(pred_mask)\n                    if annotation == \"[]\": annotation = \"authentic\"\n            \n            except Exception as e:\n                annotation = \"authentic\"\n                \n            submission_data.append({\"case_id\": str(case_id), \"annotation\": annotation})\n            \n\n    df_sub = pd.DataFrame(submission_data)\n    if not df_sub.empty:\n        df_sub = df_sub.sort_values(\"case_id\")\n        \n    df_sub.to_csv(\"submission.csv\", index=False)\n    print(\"Submission saved successfully!\")\n    \n    del model\n    del processor\n    gc.collect()\n    torch.cuda.empty_cache()\n\nif __name__ == \"__main__\":\n    if Config.TRAIN_MODE:\n        run_training()\n    else:\n        run_inference()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T08:21:14.727108Z","iopub.execute_input":"2025-12-13T08:21:14.727736Z","execution_failed":"2025-12-13T08:21:16.018Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Swin-Unet","metadata":{}},{"cell_type":"code","source":"\nimport os\nimport sys\nimport subprocess\nimport glob\nimport gc\nimport json\n\n\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\n\n\ntry:\n    import google.protobuf.message_factory\n    from google.protobuf.message_factory import MessageFactory\n    if not hasattr(MessageFactory, 'GetPrototype'):\n        def GetPrototype(self, descriptor):\n            return self.pool.GetPrototype(descriptor) if hasattr(self, 'pool') else None\n        MessageFactory.GetPrototype = GetPrototype\nexcept: pass\n\n\nprint(\"Find and load libraries online...\")\n\nsearch_path = glob.glob(\"/kaggle/input/**/my_packages\", recursive=True)\n\nif not search_path:\n    print(\"Dir my_packages nor founded\")\nelse:\n    packages_dir = search_path[0]\n    print(f\"Dir with libraries: {packages_dir}\")\n    \n    command = [\n        sys.executable, \"-m\", \"pip\", \"install\", \n        \"segmentation-models-pytorch\", \"timm\", \"efficientnet-pytorch\",\n        \"--no-index\", \n        f\"--find-links={packages_dir}\"\n    ]\n    \n    try:\n        subprocess.check_call(command)\n        print(\"SMP loaded\")\n    except subprocess.CalledProcessError as e:\n        print(f\"Error {e}\")\n\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom tqdm.auto import tqdm\n\ntry:\n    import segmentation_models_pytorch as smp\n    print(\" SMP imported with success.\")\nexcept ImportError:\n    print(\"SMP not loaded- error\")\n    sys.exit(1)\n\n\nclass Config:\n    ROOT_DIR = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n    TRAIN_IMG_AUTH = os.path.join(ROOT_DIR, \"train_images/authentic\")\n    TRAIN_IMG_FORGED = os.path.join(ROOT_DIR, \"train_images/forged\")\n    TRAIN_MASKS = os.path.join(ROOT_DIR, \"train_masks\")\n    TEST_IMGS = os.path.join(ROOT_DIR, \"test_images\")\n    \n\n    TRAIN_MODE = False\n    \n\n    SAVE_PATH = \"swin_unet_model.pth\"\n    MODEL_PATH = \"/kaggle/input/my-swin-weights/swin_unet_model.pth\"\n\n\n    ENCODER_NAME = \"tu-swin_tiny_patch4_window7_224\" \n    ENCODER_WEIGHTS = \"imagenet\"\n    \n    IMAGE_SIZE = 224\n    BATCH_SIZE = 16\n    LR = 1e-4\n    EPOCHS = 5\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\nclass ForgeryDataset(Dataset):\n    def __init__(self, image_paths, mask_dir, preprocessing_fn=None):\n        self.image_paths = image_paths\n        self.mask_dir = mask_dir\n        self.preprocessing_fn = preprocessing_fn\n        \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        file_name = os.path.basename(img_path)\n        \n        image = cv2.imread(img_path)\n        if image is None: return self._dummy()\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        orig_h, orig_w = image.shape[:2]\n        \n        mask = np.zeros((orig_h, orig_w), dtype=np.float32)\n        if 'authentic' not in img_path:\n            mask_path = os.path.join(self.mask_dir, file_name.replace('.png', '.npy'))\n            if os.path.exists(mask_path):\n                try:\n                    m = np.load(mask_path)\n                    if m.ndim == 3: m = np.max(m, axis=0) if m.shape[0] < 5 else np.max(m, axis=-1)\n                    m = np.squeeze(m)\n                    if m.shape != (orig_h, orig_w):\n                        m = cv2.resize(m.astype(np.float32), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)\n                    mask = m.astype(np.float32)\n                except: pass\n\n        image = cv2.resize(image, (Config.IMAGE_SIZE, Config.IMAGE_SIZE))\n        mask = cv2.resize(mask, (Config.IMAGE_SIZE, Config.IMAGE_SIZE), interpolation=cv2.INTER_NEAREST)\n        \n        if self.preprocessing_fn:\n            image = self.preprocessing_fn(image)\n        \n        image_tensor = torch.from_numpy(image).permute(2, 0, 1).float()\n        \n        if mask.ndim == 3: mask = mask[:, :, 0]\n        mask_tensor = torch.from_numpy(mask).unsqueeze(0).float()\n        \n        return {\n            \"image\": image_tensor,\n            \"mask\": mask_tensor,\n            \"case_id\": os.path.splitext(file_name)[0]\n        }\n\n    def _dummy(self):\n        return {\n            \"image\": torch.zeros((3, Config.IMAGE_SIZE, Config.IMAGE_SIZE)),\n            \"mask\": torch.zeros((1, Config.IMAGE_SIZE, Config.IMAGE_SIZE)),\n            \"case_id\": \"dummy\"\n        }\n\n\ndef rle_encode(mask):\n    pixels = mask.flatten(order='F')\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return json.dumps(runs.tolist())\n\n\ndef run_training():\n    print(\"TRAIN MODE: model Swin-Unet...\")\n    \n    try:\n        model = smp.Unet(\n            encoder_name=Config.ENCODER_NAME,\n            encoder_weights=Config.ENCODER_WEIGHTS, \n            in_channels=3,\n            classes=1\n        )\n    except Exception as e:\n        print(f\"Error: {e}\")\n        Config.ENCODER_NAME = \"swin_tiny_patch4_window7_224\"\n        model = smp.Unet(\n            encoder_name=Config.ENCODER_NAME,\n            encoder_weights=Config.ENCODER_WEIGHTS, \n            in_channels=3,\n            classes=1\n        )\n\n    model.to(Config.DEVICE)\n    \n    preprocessing_fn = smp.encoders.get_preprocessing_fn(Config.ENCODER_NAME, Config.ENCODER_WEIGHTS)\n    \n    bad_files = ['57085.png', '25311.png', '34265.png', '64128.png']\n    auth_files = glob.glob(os.path.join(Config.TRAIN_IMG_AUTH, \"*.png\"))\n    forg_files = [x for x in glob.glob(os.path.join(Config.TRAIN_IMG_FORGED, \"*.png\")) \n                  if os.path.basename(x) not in bad_files]\n    \n    all_files = auth_files + forg_files\n    np.random.shuffle(all_files)\n    \n    split = int(len(all_files) * 0.9)\n    train_ds = ForgeryDataset(all_files[:split], Config.TRAIN_MASKS, preprocessing_fn)\n    val_ds = ForgeryDataset(all_files[split:], Config.TRAIN_MASKS, preprocessing_fn)\n    \n    train_loader = DataLoader(train_ds, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2)\n    \n    optimizer = AdamW(model.parameters(), lr=Config.LR)\n    criterion = smp.losses.DiceLoss(mode='binary')\n    \n    print(\" Start Training...\")\n    for epoch in range(Config.EPOCHS):\n        model.train()\n        train_loss = 0\n        pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}\")\n        \n        for batch in pbar:\n            imgs = batch['image'].to(Config.DEVICE)\n            msks = batch['mask'].to(Config.DEVICE)\n            \n            optimizer.zero_grad()\n            logits = model(imgs)\n            loss = criterion(logits, msks)\n            loss.backward()\n            optimizer.step()\n            \n            train_loss += loss.item()\n            pbar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})\n            \n        print(f\"Epoch {epoch+1} Avg Loss: {train_loss/len(train_loader):.4f}\")\n        \n    torch.save(model.state_dict(), Config.SAVE_PATH)\n    print(f\"Model saved into {Config.SAVE_PATH}\")\n\n\ndef run_inference():\n    print(\"SUBMISSION MODE: Запуск инференса...\")\n    \n    if not os.path.exists(Config.MODEL_PATH):\n        print(f\"Error: file {Config.MODEL_PATH} not found!\")\n        pd.DataFrame(columns=[\"case_id\", \"annotation\"]).to_csv(\"submission.csv\", index=False)\n        return\n\n    try:\n        model = smp.Unet(\n            encoder_name=Config.ENCODER_NAME,\n            encoder_weights=None, \n            in_channels=3,\n            classes=1\n        )\n    except:\n        Config.ENCODER_NAME = \"swin_tiny_patch4_window7_224\"\n        model = smp.Unet(encoder_name=Config.ENCODER_NAME, encoder_weights=None, in_channels=3, classes=1)\n    \n    try:\n        model.load_state_dict(torch.load(Config.MODEL_PATH, map_location=Config.DEVICE))\n    except Exception as e:\n        print(f\"Error in weight: {e}\")\n        return\n\n    model.to(Config.DEVICE)\n    model.eval()\n    \n    def manual_preprocess(image):\n        image = image.astype(np.float32) / 255.0\n        image = (image - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])\n        return image\n\n    test_files = []\n    for ext in [\"*.png\", \"*.jpg\", \"*.jpeg\", \"*.tif\"]:\n        test_files.extend(glob.glob(os.path.join(Config.TEST_IMGS, ext)))\n    test_files = sorted(list(set(test_files)))\n    \n    submission_data = []\n    print(f\" Process  {len(test_files)} image...\")\n    \n    with torch.no_grad():\n        for img_path in tqdm(test_files):\n            filename = os.path.basename(img_path)\n            case_id = os.path.splitext(filename)[0]\n            \n            try:\n                image = cv2.imread(img_path)\n                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n                orig_h, orig_w = image.shape[:2]\n                \n                image_resized = cv2.resize(image, (Config.IMAGE_SIZE, Config.IMAGE_SIZE))\n                \n                # Нормализация\n                img_proc = manual_preprocess(image_resized)\n                t = torch.from_numpy(img_proc).permute(2,0,1).float().unsqueeze(0).to(Config.DEVICE)\n                \n                logits = model(t)\n                \n                logits_upsampled = nn.functional.interpolate(\n                    logits, size=(orig_h, orig_w), \n                    mode=\"bilinear\", align_corners=False\n                )\n                \n                probs = torch.sigmoid(logits_upsampled).cpu().numpy()[0, 0]\n                pred_mask = (probs > 0.5).astype(np.uint8)\n                \n                if np.sum(pred_mask) == 0:\n                    annotation = \"authentic\"\n                else:\n                    annotation = rle_encode(pred_mask)\n                    if annotation == \"[]\": annotation = \"authentic\"\n            \n            except Exception as e:\n                annotation = \"authentic\"\n                \n            submission_data.append({\"case_id\": str(case_id), \"annotation\": annotation})\n            \n    df_sub = pd.DataFrame(submission_data)\n    if not df_sub.empty:\n        df_sub = df_sub.sort_values(\"case_id\")\n        \n    df_sub.to_csv(\"submission.csv\", index=False)\n    print(\"Submission saved successfully!\")\n    \n    del model\n    gc.collect()\n    torch.cuda.empty_cache()\n\n\n#if __name__ == \"__main__\":\n#    if Config.TRAIN_MODE:\n#        run_training()\n#    else:\n#        run_inference()\n        \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T12:51:48.701142Z","iopub.execute_input":"2025-12-13T12:51:48.701788Z","iopub.status.idle":"2025-12-13T12:52:30.901829Z","shell.execute_reply.started":"2025-12-13T12:51:48.701747Z","shell.execute_reply":"2025-12-13T12:52:30.900912Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Segformer - Unet","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\nimport subprocess\nimport glob\nimport gc\nimport json\nimport random\n\n\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\n\n\ntry:\n    import google.protobuf.message_factory\n    from google.protobuf.message_factory import MessageFactory\n    if not hasattr(MessageFactory, 'GetPrototype'):\n        def GetPrototype(self, descriptor):\n            return self.pool.GetPrototype(descriptor) if hasattr(self, 'pool') else None\n        MessageFactory.GetPrototype = GetPrototype\nexcept: pass\n\nprint(\"Load library offline...\")\nsearch_path = glob.glob(\"/kaggle/input/**/my_packages\", recursive=True)\nif search_path:\n    packages_dir = search_path[0]\n    try:\n        subprocess.check_call([\n            sys.executable, \"-m\", \"pip\", \"install\", \n            \"segmentation-models-pytorch\", \"timm\", \"efficientnet-pytorch\",\n            \"--no-index\", f\"--find-links={packages_dir}\"\n        ])\n        print(\"Library downloaded!\")\n    except: pass\n\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom tqdm.auto import tqdm\nimport albumentations as A \nfrom albumentations.pytorch import ToTensorV2\n\ntry:\n    import segmentation_models_pytorch as smp\nexcept ImportError:\n    sys.exit(\"SMP not found.\")\n\n\ndef seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything()\n\n\nclass Config:\n    ROOT_DIR = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n    TRAIN_IMG_AUTH = os.path.join(ROOT_DIR, \"train_images/authentic\")\n    TRAIN_IMG_FORGED = os.path.join(ROOT_DIR, \"train_images/forged\")\n    TRAIN_MASKS = os.path.join(ROOT_DIR, \"train_masks\")\n    TEST_IMGS = os.path.join(ROOT_DIR, \"test_images\")\n    \n    TRAIN_MODE = False \n    \n    SAVE_PATH = \"best_model.pth\"\n    MODEL_PATH = \"/kaggle/input/my-segformerunet-weights/best_model(8).pth\" \n\n    ENCODER_NAME = \"mit_b3\" \n    ENCODER_WEIGHTS = \"imagenet\"\n    \n    IMAGE_SIZE = 320 \n    BATCH_SIZE = 16\n    LR = 1e-4\n    EPOCHS = 10\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndef get_train_augs():\n    return A.Compose([\n        A.Resize(Config.IMAGE_SIZE, Config.IMAGE_SIZE),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=15, p=0.4),\n        A.OneOf([\n            A.GaussNoise(var_limit=(10.0, 50.0), p=0.5),\n            A.ImageCompression(quality_lower=85, quality_upper=100, p=0.5),\n            A.CoarseDropout(max_holes=8, max_height=16, max_width=16, p=0.2),\n        ], p=0.3),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ])\n\ndef get_valid_augs():\n    return A.Compose([\n        A.Resize(Config.IMAGE_SIZE, Config.IMAGE_SIZE),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ])\n\n\nclass ForgeryDataset(Dataset):\n    def __init__(self, image_paths, mask_dir, transforms=None):\n        self.image_paths = image_paths\n        self.mask_dir = mask_dir\n        self.transforms = transforms\n        \n    def __len__(self): return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        file_name = os.path.basename(img_path)\n        image = cv2.imread(img_path)\n        if image is None: return self._dummy()\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        orig_h, orig_w = image.shape[:2]\n        \n        mask = np.zeros((orig_h, orig_w), dtype=np.float32)\n        if 'authentic' not in img_path:\n            mask_path = os.path.join(self.mask_dir, file_name.replace('.png', '.npy'))\n            if os.path.exists(mask_path):\n                try:\n                    m = np.load(mask_path)\n                    if m.ndim == 3: m = np.max(m, axis=0) if m.shape[0] < 5 else np.max(m, axis=-1)\n                    m = np.squeeze(m)\n                    if m.shape != (orig_h, orig_w):\n                        m = cv2.resize(m.astype(np.float32), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)\n                    mask = m.astype(np.float32)\n                except: pass\n\n        if self.transforms:\n            augmented = self.transforms(image=image, mask=mask)\n            image = augmented['image']\n            mask = augmented['mask'].unsqueeze(0).float()\n        else:\n            image = cv2.resize(image, (Config.IMAGE_SIZE, Config.IMAGE_SIZE))\n            image = torch.from_numpy(image).permute(2,0,1).float() / 255.0\n            mask = cv2.resize(mask, (Config.IMAGE_SIZE, Config.IMAGE_SIZE), interpolation=cv2.INTER_NEAREST)\n            mask = torch.from_numpy(mask).unsqueeze(0).float()\n        \n        return {\"image\": image, \"mask\": mask, \"case_id\": os.path.splitext(file_name)[0]}\n\n    def _dummy(self):\n        return {\n            \"image\": torch.zeros((3, Config.IMAGE_SIZE, Config.IMAGE_SIZE)),\n            \"mask\": torch.zeros((1, Config.IMAGE_SIZE, Config.IMAGE_SIZE)),\n            \"case_id\": \"dummy\"\n        }\n\ndef rle_encode(mask):\n    pixels = mask.flatten(order='F')\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return json.dumps(runs.tolist())\n\n\ndef run_training():\n    print(f\" TRAIN MODE:  {Config.ENCODER_NAME}...\")\n    \n    model = smp.Unet(\n        encoder_name=Config.ENCODER_NAME,\n        encoder_weights=Config.ENCODER_WEIGHTS, \n        in_channels=3, classes=1,\n        decoder_channels=(256, 128, 64, 32, 16) \n    )\n    model.to(Config.DEVICE)\n    \n    bad_files = ['57085.png', '25311.png', '34265.png', '64128.png']\n    auth_files = glob.glob(os.path.join(Config.TRAIN_IMG_AUTH, \"*.png\"))\n    forg_files = [x for x in glob.glob(os.path.join(Config.TRAIN_IMG_FORGED, \"*.png\")) \n                  if os.path.basename(x) not in bad_files]\n    \n    all_files = auth_files + forg_files\n    np.random.shuffle(all_files)\n    \n    split = int(len(all_files) * 0.9)\n    train_ds = ForgeryDataset(all_files[:split], Config.TRAIN_MASKS, get_train_augs())\n    val_ds = ForgeryDataset(all_files[split:], Config.TRAIN_MASKS, get_valid_augs())\n    \n    train_loader = DataLoader(train_ds, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2)\n    val_loader = DataLoader(val_ds, batch_size=Config.BATCH_SIZE, shuffle=False, num_workers=2)\n    \n    optimizer = AdamW(model.parameters(), lr=Config.LR, weight_decay=1e-3)\n    scheduler = CosineAnnealingLR(optimizer, T_max=Config.EPOCHS, eta_min=1e-6)\n    \n    dice_loss = smp.losses.DiceLoss(mode='binary')\n    bce_loss = nn.BCEWithLogitsLoss()\n    def criterion(pred, target): return 0.5*dice_loss(pred, target) + 0.5*bce_loss(pred, target)\n    \n    best_loss = float('inf')\n    \n    print(\"Start Training...\")\n    for epoch in range(Config.EPOCHS):\n        model.train()\n        train_loss = 0\n        pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}\")\n        \n        for batch in pbar:\n            imgs = batch['image'].to(Config.DEVICE)\n            msks = batch['mask'].to(Config.DEVICE)\n            optimizer.zero_grad()\n            logits = model(imgs)\n            loss = criterion(logits, msks)\n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item()\n            pbar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})\n        \n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for batch in val_loader:\n                imgs = batch['image'].to(Config.DEVICE)\n                msks = batch['mask'].to(Config.DEVICE)\n                logits = model(imgs)\n                val_loss += criterion(logits, msks).item()\n        \n        avg_val = val_loss / len(val_loader)\n        scheduler.step()\n        print(f\"Epoch {epoch+1}: Val Loss={avg_val:.4f}\")\n        \n        if avg_val < best_loss:\n            best_loss = avg_val\n            torch.save(model.state_dict(), Config.SAVE_PATH)\n            print(f\"Saved Best Model!\")\n            \n    print(\"Learning completed.\")\n\ndef run_inference():\n    print(\" SUBMISSION MODE...\")\n    \n    if not os.path.exists(Config.MODEL_PATH):\n        print(f\"Fail {Config.MODEL_PATH} not found.\")\n        pd.DataFrame(columns=[\"case_id\", \"annotation\"]).to_csv(\"submission.csv\", index=False)\n        return\n\n    try:\n        model = smp.Unet(\n            encoder_name=Config.ENCODER_NAME,\n            encoder_weights=None, \n            in_channels=3, classes=1,\n            decoder_channels=(256, 128, 64, 32, 16)\n        )\n        model.load_state_dict(torch.load(Config.MODEL_PATH, map_location=Config.DEVICE))\n    except Exception as e:\n        print(f\"Error: {e}\")\n        return\n\n    model.to(Config.DEVICE)\n    model.eval()\n    \n    transforms = get_valid_augs()\n    test_files = sorted(glob.glob(os.path.join(Config.TEST_IMGS, \"*\")))\n    submission_data = []\n    \n    print(f\"Process {len(test_files)} image...\")\n    with torch.no_grad():\n        for img_path in tqdm(test_files):\n            case_id = os.path.splitext(os.path.basename(img_path))[0]\n            try:\n                image = cv2.imread(img_path)\n                orig_h, orig_w = image.shape[:2]\n                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n                \n                augmented = transforms(image=image)\n                t = augmented['image'].unsqueeze(0).to(Config.DEVICE)\n                \n                logits = model(t)\n                logits_up = nn.functional.interpolate(logits, size=(orig_h, orig_w), mode=\"bilinear\", align_corners=False)\n                pred_mask = (torch.sigmoid(logits_up) > 0.5).cpu().numpy()[0,0].astype(np.uint8)\n                \n                if np.sum(pred_mask) == 0: annotation = \"authentic\"\n                else: \n                    annotation = rle_encode(pred_mask)\n                    if annotation == \"[]\": annotation = \"authentic\"\n            except: annotation = \"authentic\"\n            submission_data.append({\"case_id\": str(case_id), \"annotation\": annotation})\n            \n    df = pd.DataFrame(submission_data).sort_values(\"case_id\")\n    df.to_csv(\"submission.csv\", index=False)\n    print(\"Submission saved!\")\n    del model\n    gc.collect()\n    torch.cuda.empty_cache()\n\n# if __name__ == \"__main__\":\n#     if Config.TRAIN_MODE:\n#         run_training()\n#     else:\n#         run_inference()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}