{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":19989,"databundleVersionId":1160143},{"sourceType":"datasetVersion","sourceId":2184410,"datasetId":1311358,"databundleVersionId":2225813},{"sourceType":"datasetVersion","sourceId":492536,"datasetId":230545,"databundleVersionId":508605}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install scipy tqdm -q\n\nimport numpy as np\nimport pandas as pd\nimport json\nimport cv2\nimport os\nfrom scipy.io import loadmat\nfrom scipy.spatial import KDTree\nimport matplotlib.pyplot as plt\n\nprint(\"基础库导入完成\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:52:56.757459Z","iopub.execute_input":"2026-04-11T14:52:56.758161Z","iopub.status.idle":"2026-04-11T14:53:00.024479Z","shell.execute_reply.started":"2026-04-11T14:52:56.758128Z","shell.execute_reply":"2026-04-11T14:53:00.023592Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DensityMapGenerator:\n    def __init__(self, sigma=15, beta=0.3):\n        self.sigma = sigma\n        self.beta = beta\n    \n    def gaussian_kernel(self, size, sigma):\n        kernel = np.zeros((size, size))\n        center = size // 2\n        for i in range(size):\n            for j in range(size):\n                x = i - center\n                y = j - center\n                kernel[i, j] = np.exp(-(x**2 + y**2) / (2 * sigma**2))\n        return kernel / kernel.sum()\n    \n    def geometry_adaptive_kernel(self, points, beta=0.3):\n        if len(points) <= 1:\n            return [self.sigma] * len(points)\n        tree = KDTree(points)\n        k = min(3, len(points))\n        distances, _ = tree.query(points, k=k+1)\n        sigmas = []\n        for i, dist in enumerate(distances):\n            if k == 1:\n                avg_dist = self.sigma\n            else:\n                avg_dist = np.mean(dist[1:])\n            sigmas.append(avg_dist * beta)\n        return sigmas\n    \n    def from_points(self, img_shape, points, method='fixed', beta=0.3):\n        h, w = img_shape[:2]\n        density_map = np.zeros((h, w), dtype=np.float32)\n        if len(points) == 0:\n            return density_map\n        \n        if method == 'adaptive':\n            sigmas = self.geometry_adaptive_kernel(points, beta)\n        else:\n            sigmas = [self.sigma] * len(points)\n        \n        for (x, y), sigma in zip(points, sigmas):\n            if 0 <= x < w and 0 <= y < h:\n                kernel_size = int(6 * sigma)\n                if kernel_size % 2 == 0:\n                    kernel_size += 1\n                kernel_size = max(3, min(kernel_size, 99))\n                kernel = self.gaussian_kernel(kernel_size, sigma)\n                \n                x_start = int(x - kernel_size // 2)\n                y_start = int(y - kernel_size // 2)\n                x_end = x_start + kernel_size\n                y_end = y_start + kernel_size\n                \n                k_x_start = max(0, -x_start)\n                k_y_start = max(0, -y_start)\n                k_x_end = kernel_size - max(0, x_end - w)\n                k_y_end = kernel_size - max(0, y_end - h)\n                \n                x_start = max(0, x_start)\n                y_start = max(0, y_start)\n                x_end = min(w, x_end)\n                y_end = min(h, y_end)\n                \n                density_map[y_start:y_end, x_start:x_end] += \\\n                    kernel[k_y_start:k_y_end, k_x_start:k_x_end]\n        return density_map\n    \n    def from_bboxes(self, img_shape, bboxes, method='center'):\n        if method == 'center':\n            points = []\n            for bbox in bboxes:\n                if len(bbox) == 4:\n                    if bbox[2] > bbox[0] and bbox[3] > bbox[1]:\n                        cx = (bbox[0] + bbox[2]) / 2\n                        cy = (bbox[1] + bbox[3]) / 2\n                    else:\n                        cx = bbox[0] + bbox[2] / 2\n                        cy = bbox[1] + bbox[3] / 2\n                    points.append([cx, cy])\n            return self.from_points(img_shape, points, method='fixed')\n        else:\n            h, w = img_shape[:2]\n            density_map = np.zeros((h, w), dtype=np.float32)\n            for bbox in bboxes:\n                if len(bbox) == 4:\n                    if bbox[2] > bbox[0] and bbox[3] > bbox[1]:\n                        xmin, ymin, xmax, ymax = bbox\n                        box_w, box_h = xmax - xmin, ymax - ymin\n                    else:\n                        x, y, box_w, box_h = bbox\n                        xmin, ymin, xmax, ymax = x, y, x+box_w, y+box_h\n                    sigma = max(box_w, box_h) / 4\n                    kernel_size = int(6 * sigma)\n                    if kernel_size % 2 == 0:\n                        kernel_size += 1\n                    kernel_size = max(3, min(kernel_size, 99))\n                    kernel = self.gaussian_kernel(kernel_size, sigma)\n                    cx = (xmin + xmax) / 2\n                    cy = (ymin + ymax) / 2\n                    x_start = int(cx - kernel_size // 2)\n                    y_start = int(cy - kernel_size // 2)\n                    x_end = x_start + kernel_size\n                    y_end = y_start + kernel_size\n                    k_x_start = max(0, -x_start)\n                    k_y_start = max(0, -y_start)\n                    k_x_end = kernel_size - max(0, x_end - w)\n                    k_y_end = kernel_size - max(0, y_end - h)\n                    x_start = max(0, x_start)\n                    y_start = max(0, y_start)\n                    x_end = min(w, x_end)\n                    y_end = min(h, y_end)\n                    density_map[y_start:y_end, x_start:x_end] += \\\n                        kernel[k_y_start:k_y_end, k_x_start:k_x_end]\n            return density_map\n\nprint(\"DensityMapGenerator 定义完成\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.026487Z","iopub.execute_input":"2026-04-11T14:53:00.026735Z","iopub.status.idle":"2026-04-11T14:53:00.045194Z","shell.execute_reply.started":"2026-04-11T14:53:00.026708Z","shell.execute_reply":"2026-04-11T14:53:00.044474Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.optim as optim\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nimport torchvision.models as models\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"PyTorch版本: {torch.__version__}\")\nprint(f\"CUDA可用: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"GPU型号: {torch.cuda.get_device_name(0)}\")\n    print(f\"GPU数量: {torch.cuda.device_count()}\")\nprint(f\"使用设备: {device}\")\n\n# 密度图放大因子（关键修复）\nDENSITY_SCALE = 100.0\nprint(f\"密度图放大因子: {DENSITY_SCALE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.045904Z","iopub.execute_input":"2026-04-11T14:53:00.046092Z","iopub.status.idle":"2026-04-11T14:53:00.062192Z","shell.execute_reply.started":"2026-04-11T14:53:00.046074Z","shell.execute_reply":"2026-04-11T14:53:00.061497Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MultiDatasetCounter(Dataset):\n    \"\"\"统一的多数据集计数数据加载器（修复版）\"\"\"\n    def __init__(self, dataset_type='shanghai', split='train', part='A', \n                 transform=None, target_size=(512, 512)):\n        self.dataset_type = dataset_type\n        self.split = split\n        self.part = part\n        self.transform = transform\n        self.target_size = target_size\n        self.output_size = (target_size[0] // 8, target_size[1] // 8)\n        \n        # 数据集特定配置\n        self.sigma_config = {\n            'shanghai': {'sigma': 8, 'method': 'adaptive', 'beta': 0.3},\n            'bccd': {'sigma': 15, 'method': 'fixed', 'beta': 0.3},\n            'wheat': {'sigma': 6, 'method': 'fixed', 'beta': 0.3}\n        }\n        \n        self.paths = {\n            'shanghai': {\n                'part_A': '/kaggle/input/datasets/tthien/shanghaitech/ShanghaiTech/part_A',\n                'part_B': '/kaggle/input/datasets/tthien/shanghaitech/ShanghaiTech/part_B'\n            },\n            'bccd': {\n                'train': '/kaggle/input/datasets/ammarnassanalhajali/bccd-coco/train',\n                'valid': '/kaggle/input/datasets/ammarnassanalhajali/bccd-coco/valid',\n                'test': '/kaggle/input/datasets/ammarnassanalhajali/bccd-coco/test'\n            },\n            'wheat': {\n                'images': '/kaggle/input/competitions/global-wheat-detection/train',\n                'csv': '/kaggle/input/competitions/global-wheat-detection/train.csv'\n            }\n        }\n        \n        self.data_list = self._load_data_list()\n        \n    def _load_data_list(self):\n        data_list = []\n        \n        if self.dataset_type == 'shanghai':\n            data_path = os.path.join(self.paths['shanghai'][f'part_{self.part}'], f'{self.split}_data')\n            img_dir = os.path.join(data_path, 'images')\n            gt_dir = os.path.join(data_path, 'ground-truth')\n            \n            img_files = sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg')])\n            for img_file in img_files:\n                img_path = os.path.join(img_dir, img_file)\n                mat_file = f\"GT_{img_file.replace('.jpg', '.mat')}\"\n                gt_path = os.path.join(gt_dir, mat_file)\n                data_list.append({'img': img_path, 'gt': gt_path, 'type': 'shanghai'})\n                \n        elif self.dataset_type == 'bccd':\n            data_path = self.paths['bccd'][self.split]\n            json_path = os.path.join(data_path, '_annotations.coco.json')\n            \n            with open(json_path, 'r') as f:\n                coco_data = json.load(f)\n            \n            id_to_file = {img['id']: img['file_name'] for img in coco_data['images']}\n            img_annotations = {}\n            for ann in coco_data['annotations']:\n                img_id = ann['image_id']\n                if img_id not in img_annotations:\n                    img_annotations[img_id] = []\n                img_annotations[img_id].append(ann['bbox'])\n            \n            for img_id, bboxes in img_annotations.items():\n                img_path = os.path.join(data_path, id_to_file[img_id])\n                data_list.append({'img': img_path, 'bboxes': bboxes, 'type': 'bccd'})\n                \n        elif self.dataset_type == 'wheat':\n            df = pd.read_csv(self.paths['wheat']['csv'])\n            for img_id in df['image_id'].unique():\n                img_path = os.path.join(self.paths['wheat']['images'], f\"{img_id}.jpg\")\n                if not os.path.exists(img_path):\n                    continue\n                img_data = df[df['image_id'] == img_id]\n                bboxes = []\n                for _, row in img_data.iterrows():\n                    bbox_str = row['bbox']\n                    bbox = eval(bbox_str) if isinstance(bbox_str, str) else bbox_str\n                    bboxes.append(bbox)\n                data_list.append({'img': img_path, 'bboxes': bboxes, 'type': 'wheat'})\n        \n        return data_list\n    \n    def __len__(self):\n        return len(self.data_list)\n    \n    def __getitem__(self, idx):\n        item = self.data_list[idx]\n        \n        cfg = self.sigma_config[self.dataset_type]\n        generator = DensityMapGenerator(sigma=cfg['sigma'], beta=cfg['beta'])\n        \n        img = cv2.imread(item['img'])\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        # 生成密度图\n        if item['type'] == 'shanghai':\n            mat_data = loadmat(item['gt'])\n            points = mat_data['image_info'][0][0][0][0][0]\n            if points.shape[1] == 2:\n                points = points.tolist()\n            else:\n                points = points.T.tolist()\n            density = generator.from_points(img.shape, points, method=cfg['method'], beta=cfg['beta'])\n            count = len(points)\n        else:\n            density = generator.from_bboxes(img.shape, item['bboxes'], method='center')\n            count = len(item['bboxes'])\n        \n        # 调整图像\n        img = cv2.resize(img, self.target_size, interpolation=cv2.INTER_LINEAR)\n        \n        # 调整密度图\n        density = cv2.resize(density, self.output_size, interpolation=cv2.INTER_LINEAR)\n        \n        # ========== 关键修复 ==========\n        # 1. 用真实计数归一化\n        density_sum = density.sum()\n        if density_sum > 0:\n            density = density * (count / density_sum)\n        elif count > 0:\n            density = np.ones(self.output_size, dtype=np.float32) * (count / (self.output_size[0] * self.output_size[1]))\n        \n        # 2. 放大密度图，增加动态范围\n        density = density * DENSITY_SCALE\n        # ==============================\n        \n        img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0\n        density = torch.from_numpy(density).unsqueeze(0).float()\n        \n        return img, density, count\n\nprint(\"MultiDatasetCounter 定义完成\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.063675Z","iopub.execute_input":"2026-04-11T14:53:00.063890Z","iopub.status.idle":"2026-04-11T14:53:00.081433Z","shell.execute_reply.started":"2026-04-11T14:53:00.063871Z","shell.execute_reply":"2026-04-11T14:53:00.080845Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CSRNet(nn.Module):\n    \"\"\"CSRNet: Dilated Convolutional Neural Networks for Crowd Counting\"\"\"\n    def __init__(self, load_weights=True):\n        super(CSRNet, self).__init__()\n        \n        self.frontend = self._make_frontend()\n        self.backend = self._make_backend()\n        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)\n        \n        self._initialize_weights()\n        \n        if load_weights:\n            self._load_pretrained_vgg()\n    \n    def _make_frontend(self):\n        layers = []\n        in_channels = 3\n        cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]\n        \n        for v in cfg:\n            if v == 'M':\n                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))\n            else:\n                layers.append(nn.Conv2d(in_channels, v, kernel_size=3, padding=1))\n                layers.append(nn.ReLU(inplace=True))\n                in_channels = v\n        \n        return nn.Sequential(*layers)\n    \n    def _make_backend(self):\n        layers = []\n        dilation_rates = [2, 2, 2, 4, 4, 4]\n        in_channels = 512\n        \n        for dilation in dilation_rates:\n            layers.append(nn.Conv2d(in_channels, 512, kernel_size=3, padding=dilation, dilation=dilation))\n            layers.append(nn.ReLU(inplace=True))\n            in_channels = 512\n        \n        layers.append(nn.Conv2d(512, 256, kernel_size=3, padding=1))\n        layers.append(nn.ReLU(inplace=True))\n        layers.append(nn.Conv2d(256, 128, kernel_size=3, padding=1))\n        layers.append(nn.ReLU(inplace=True))\n        layers.append(nn.Conv2d(128, 64, kernel_size=3, padding=1))\n        layers.append(nn.ReLU(inplace=True))\n        \n        return nn.Sequential(*layers)\n    \n    def _initialize_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.normal_(m.weight, std=0.01)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n    \n    def _load_pretrained_vgg(self):\n        try:\n            vgg16 = models.vgg16(pretrained=True)\n            frontend_dict = self.frontend.state_dict()\n            vgg_dict = vgg16.features.state_dict()\n            \n            frontend_keys = list(frontend_dict.keys())\n            \n            for i in range(0, len(frontend_keys), 2):\n                if i < 20:\n                    if frontend_keys[i] in vgg_dict:\n                        frontend_dict[frontend_keys[i]] = vgg_dict[frontend_keys[i]]\n            \n            self.frontend.load_state_dict(frontend_dict, strict=False)\n            print(\"✅ VGG16预训练权重加载成功\")\n        except Exception as e:\n            print(f\"⚠️ VGG16权重加载失败: {e}\")\n    \n    def forward(self, x):\n        x = self.frontend(x)\n        x = self.backend(x)\n        x = self.output_layer(x)\n        return x\n\nprint(\"CSRNet 定义完成\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.082178Z","iopub.execute_input":"2026-04-11T14:53:00.082375Z","iopub.status.idle":"2026-04-11T14:53:00.103905Z","shell.execute_reply.started":"2026-04-11T14:53:00.082356Z","shell.execute_reply":"2026-04-11T14:53:00.103226Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CSRNetTrainer:\n    def __init__(self, model, device, learning_rate=1e-5):\n        self.model = model.to(device)\n        self.device = device\n        self.optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.95, weight_decay=5e-4)\n        self.scheduler = ReduceLROnPlateau(self.optimizer, mode='min', factor=0.5, patience=5)\n        self.criterion = nn.MSELoss()\n        \n        self.train_losses = []\n        self.val_losses = []\n        self.best_mae = float('inf')\n        \n    def train_epoch(self, train_loader):\n        self.model.train()\n        epoch_loss = 0\n        \n        pbar = tqdm(train_loader, desc='Training')\n        for imgs, densities, counts in pbar:\n            imgs = imgs.to(self.device, non_blocking=True)\n            densities = densities.to(self.device, non_blocking=True)\n            \n            self.optimizer.zero_grad()\n            outputs = self.model(imgs)\n            loss = self.criterion(outputs, densities)\n            loss.backward()\n            self.optimizer.step()\n            \n            epoch_loss += loss.item()\n            pbar.set_postfix({'Loss': f'{loss.item():.4f}'})\n        \n        return epoch_loss / len(train_loader)\n    \n    def validate(self, val_loader):\n        self.model.eval()\n        val_loss = 0\n        mae = 0\n        mse = 0\n        total_samples = 0\n        \n        with torch.no_grad():\n            for imgs, densities, counts in tqdm(val_loader, desc='Validating'):\n                imgs = imgs.to(self.device, non_blocking=True)\n                densities = densities.to(self.device, non_blocking=True)\n                \n                outputs = self.model(imgs)\n                loss = self.criterion(outputs, densities)\n                val_loss += loss.item()\n                \n                # 预测计数需要除以放大因子\n                pred_counts = outputs.sum(dim=(1, 2, 3)).cpu().numpy() / DENSITY_SCALE\n                true_counts = counts.numpy()\n                \n                mae += np.abs(pred_counts - true_counts).sum()\n                mse += ((pred_counts - true_counts) ** 2).sum()\n                total_samples += imgs.size(0)\n        \n        avg_loss = val_loss / len(val_loader)\n        avg_mae = mae / total_samples\n        avg_mse = np.sqrt(mse / total_samples)\n        \n        return avg_loss, avg_mae, avg_mse\n    \n    def train(self, train_loader, val_loader, epochs, save_path='./checkpoints_csrnet'):\n        os.makedirs(save_path, exist_ok=True)\n        \n        for epoch in range(epochs):\n            print(f'\\nEpoch {epoch+1}/{epochs}')\n            print('-' * 50)\n            \n            train_loss = self.train_epoch(train_loader)\n            self.train_losses.append(train_loss)\n            \n            val_loss, val_mae, val_mse = self.validate(val_loader)\n            self.val_losses.append(val_loss)\n            \n            self.scheduler.step(val_mae)\n            current_lr = self.optimizer.param_groups[0]['lr']\n            \n            print(f'Train Loss: {train_loss:.4f}')\n            print(f'Val Loss: {val_loss:.4f}, MAE: {val_mae:.2f}, MSE: {val_mse:.2f}')\n            print(f'Learning Rate: {current_lr:.6f}')\n            \n            if torch.cuda.is_available():\n                print(f'GPU Memory: {torch.cuda.memory_allocated()/1e9:.2f}/{torch.cuda.memory_reserved()/1e9:.2f} GB')\n            \n            if val_mae < self.best_mae:\n                self.best_mae = val_mae\n                if isinstance(self.model, nn.DataParallel):\n                    model_state = self.model.module.state_dict()\n                else:\n                    model_state = self.model.state_dict()\n                    \n                torch.save({\n                    'epoch': epoch,\n                    'model_state_dict': model_state,\n                    'optimizer_state_dict': self.optimizer.state_dict(),\n                    'best_mae': self.best_mae,\n                }, os.path.join(save_path, 'best_model.pth'))\n                print(f'✓ Saved best model with MAE: {val_mae:.2f}')\n        \n        return self.train_losses, self.val_losses\n\nprint(\"CSRNetTrainer 定义完成\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.104946Z","iopub.execute_input":"2026-04-11T14:53:00.105296Z","iopub.status.idle":"2026-04-11T14:53:00.120123Z","shell.execute_reply.started":"2026-04-11T14:53:00.105253Z","shell.execute_reply":"2026-04-11T14:53:00.119410Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Visualizer:\n    @staticmethod\n    def predict_and_visualize(model, dataset, device, num_samples=3, save_name='samples'):\n        model.eval()\n        \n        if isinstance(dataset, torch.utils.data.Subset):\n            indices = np.random.choice(len(dataset), min(num_samples, len(dataset)), replace=False)\n        else:\n            indices = np.random.choice(len(dataset), min(num_samples, len(dataset)), replace=False)\n        \n        fig, axes = plt.subplots(num_samples, 4, figsize=(16, 4*num_samples))\n        if num_samples == 1:\n            axes = axes.reshape(1, -1)\n        \n        for i, idx in enumerate(indices):\n            img, density, count = dataset[idx]\n            \n            with torch.no_grad():\n                img_tensor = img.unsqueeze(0).to(device)\n                pred_density = model(img_tensor).squeeze().cpu().numpy()\n            \n            img_np = img.permute(1, 2, 0).numpy()\n            \n            axes[i, 0].imshow(img_np)\n            axes[i, 0].set_title(f'Original\\nTrue: {count}')\n            axes[i, 0].axis('off')\n            \n            true_density = density.squeeze().numpy()\n            true_disp = cv2.resize(true_density, img_np.shape[:2][::-1])\n            im1 = axes[i, 1].imshow(true_disp, cmap='jet')\n            axes[i, 1].set_title(f'GT Density\\nSum: {true_density.sum()/DENSITY_SCALE:.1f}')\n            axes[i, 1].axis('off')\n            plt.colorbar(im1, ax=axes[i, 1], fraction=0.046)\n            \n            pred_disp = cv2.resize(pred_density, img_np.shape[:2][::-1])\n            im2 = axes[i, 2].imshow(pred_disp, cmap='jet')\n            pred_count = pred_density.sum() / DENSITY_SCALE\n            axes[i, 2].set_title(f'Pred Density\\nSum: {pred_count:.1f}')\n            axes[i, 2].axis('off')\n            plt.colorbar(im2, ax=axes[i, 2], fraction=0.046)\n            \n            error = abs(pred_count - count) / count * 100 if count > 0 else 0\n            axes[i, 3].bar(['True', 'Pred'], [count, pred_count], color=['#2E86AB', '#A23B72'])\n            axes[i, 3].set_title(f'Count Comparison\\nError: {error:.1f}%')\n            axes[i, 3].set_ylabel('Count')\n            \n            for j, v in enumerate([count, pred_count]):\n                axes[i, 3].text(j, v + max(count, pred_count)*0.02, f'{v:.1f}', ha='center', va='bottom')\n        \n        plt.tight_layout()\n        plt.savefig(f'/kaggle/working/{save_name}.png', dpi=150, bbox_inches='tight')\n        plt.show()\n    \n    @staticmethod\n    def plot_training_curves(train_losses, val_losses, save_name='training_curves'):\n        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n        epochs = range(1, len(train_losses) + 1)\n        \n        ax1.plot(epochs, train_losses, 'b-', label='Training Loss', linewidth=2)\n        ax1.plot(epochs, val_losses, 'r-', label='Validation Loss', linewidth=2)\n        ax1.set_xlabel('Epochs')\n        ax1.set_ylabel('Loss')\n        ax1.set_title('Training and Validation Loss')\n        ax1.legend()\n        ax1.grid(True, alpha=0.3)\n        \n        ax2.plot(epochs, train_losses, 'b-', alpha=0.7, label='Train')\n        ax2.plot(epochs, val_losses, 'r-', alpha=0.7, label='Validation')\n        ax2.set_xlabel('Epochs')\n        ax2.set_ylabel('Loss (Log Scale)')\n        ax2.set_title('Loss Curves (Log Scale)')\n        ax2.set_yscale('log')\n        ax2.legend()\n        ax2.grid(True, alpha=0.3)\n        \n        plt.tight_layout()\n        plt.savefig(f'/kaggle/working/{save_name}.png', dpi=150, bbox_inches='tight')\n        plt.show()\n\nprint(\"Visualizer 定义完成\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.121059Z","iopub.execute_input":"2026-04-11T14:53:00.121665Z","iopub.status.idle":"2026-04-11T14:53:00.140484Z","shell.execute_reply.started":"2026-04-11T14:53:00.121631Z","shell.execute_reply":"2026-04-11T14:53:00.139988Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n    config = {\n        'batch_size': 8,\n        'epochs': 25,\n        'learning_rate': 1e-4,\n        'image_size': (512, 512),\n        'num_workers': 4\n    }\n    \n    datasets_info = [\n        #{'type': 'shanghai', 'part': 'A'},\n        #{'type': 'bccd', 'split': 'train'},\n        {'type': 'wheat'}\n    ]\n    \n    results = {}\n    \n    for dataset_info in datasets_info:\n        dataset_name = dataset_info['type']\n        print(f'\\n{\"=\"*60}')\n        print(f'Training CSRNet on {dataset_name.upper()} Dataset')\n        print(\"=\"*60)\n        \n        if dataset_name == 'shanghai':\n            train_dataset = MultiDatasetCounter(dataset_type='shanghai', split='train', part='A', target_size=config['image_size'])\n            val_dataset = MultiDatasetCounter(dataset_type='shanghai', split='test', part='A', target_size=config['image_size'])\n        elif dataset_name == 'bccd':\n            train_dataset = MultiDatasetCounter(dataset_type='bccd', split='train', target_size=config['image_size'])\n            val_dataset = MultiDatasetCounter(dataset_type='bccd', split='valid', target_size=config['image_size'])\n        else:\n            full_dataset = MultiDatasetCounter(dataset_type='wheat', target_size=config['image_size'])\n            train_size = int(0.8 * len(full_dataset))\n            val_size = len(full_dataset) - train_size\n            train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size])\n        \n        train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], pin_memory=True)\n        val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], pin_memory=True)\n        \n        print(f'Train samples: {len(train_dataset)}')\n        print(f'Validation samples: {len(val_dataset)}')\n        \n        model = CSRNet(load_weights=True)\n        model = model.to(device)\n        \n        if torch.cuda.device_count() > 1:\n            model = nn.DataParallel(model)\n            print(f\"使用 {torch.cuda.device_count()} 个GPU\")\n        \n        trainer = CSRNetTrainer(model, device, learning_rate=config['learning_rate'])\n        \n        save_path = f'/kaggle/working/csrnet_checkpoints_fixed/{dataset_name}'\n        train_losses, val_losses = trainer.train(train_loader, val_loader, epochs=config['epochs'], save_path=save_path)\n        \n        Visualizer.plot_training_curves(train_losses, val_losses, save_name=f'{dataset_name}_curves_fixed')\n        \n        checkpoint = torch.load(os.path.join(save_path, 'best_model.pth'), map_location=device)\n        if isinstance(model, nn.DataParallel):\n            model.module.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            model.load_state_dict(checkpoint['model_state_dict'])\n        model = model.to(device)\n        \n        print(f'\\n✅ Best MAE on {dataset_name}: {checkpoint[\"best_mae\"]:.2f}')\n        \n        Visualizer.predict_and_visualize(model, val_dataset, device, num_samples=3, save_name=f'{dataset_name}_predictions_fixed')\n        \n        results[dataset_name] = {'best_mae': checkpoint['best_mae'], 'train_losses': train_losses, 'val_losses': val_losses}\n        \n        torch.cuda.empty_cache()\n    \n    print('\\n' + '='*60)\n    print('FINAL RESULTS - CSRNet (Fixed)')\n    print('='*60)\n    for dataset_name, result in results.items():\n        print(f'{dataset_name.upper()}: Best MAE = {result[\"best_mae\"]:.2f}')\n    \n    return results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.141142Z","iopub.execute_input":"2026-04-11T14:53:00.141391Z","iopub.status.idle":"2026-04-11T14:53:00.154899Z","shell.execute_reply.started":"2026-04-11T14:53:00.141371Z","shell.execute_reply":"2026-04-11T14:53:00.154318Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    print(\"\\n\" + \"=\"*60)\n    print(\"CSRNet Training - Fixed Version\")\n    print(f\"Density Scale Factor: {DENSITY_SCALE}\")\n    print(f\"Device: {device}\")\n    print(\"=\"*60 + \"\\n\")\n    \n    results = main()\n    \n    print(\"\\n\" + \"=\"*60)\n    print(\"✅ 训练完成！\")\n    print(\"=\"*60)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-11T14:53:00.155651Z","iopub.execute_input":"2026-04-11T14:53:00.155904Z","iopub.status.idle":"2026-04-11T15:27:11.619737Z","shell.execute_reply.started":"2026-04-11T14:53:00.155876Z","shell.execute_reply":"2026-04-11T15:27:11.618436Z"}},"outputs":[],"execution_count":null}]}