{"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":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip -q install -U diffusers transformers accelerate safetensors huggingface_hub xformers","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-31T06:43:11.426826Z","iopub.execute_input":"2025-12-31T06:43:11.426998Z","iopub.status.idle":"2025-12-31T06:46:56.470729Z","shell.execute_reply.started":"2025-12-31T06:43:11.426982Z","shell.execute_reply":"2025-12-31T06:46:56.469772Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom diffusers import StableDiffusionXLPipeline","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:02:17.615766Z","iopub.execute_input":"2025-12-31T10:02:17.616387Z","iopub.status.idle":"2025-12-31T10:02:51.919342Z","shell.execute_reply.started":"2025-12-31T10:02:17.616357Z","shell.execute_reply":"2025-12-31T10:02:51.918663Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nos.makedirs(\"/kaggle/working/hf_home\",exist_ok=True)\nos.environ[\"HF_HOME\"] = \"/kaggle/working/hf_home\"\nos.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"  # 可选：有时更快\n\n# 从 Kaggle Secrets 读取\nhf_token = os.environ.get(\"HF_TOKEN\")\nprint(\"HF_TOKEN exists:\", hf_token is not None)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:02:51.920866Z","iopub.execute_input":"2025-12-31T10:02:51.921583Z","iopub.status.idle":"2025-12-31T10:02:51.928445Z","shell.execute_reply.started":"2025-12-31T10:02:51.921552Z","shell.execute_reply":"2025-12-31T10:02:51.927673Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 初步测试","metadata":{}},{"cell_type":"code","source":"# model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\n\n# pipe = StableDiffusionXLPipeline.from_pretrained(\n#     model_id,\n#     torch_dtype=torch.float16,\n#     use_safetensors=True,\n#     token=hf_token,  # 需要权限/同意条款时很关键\n# ).to(\"cuda\")\n\n# # 省显存/更稳（尤其是 T4）\n# pipe.enable_vae_slicing()\n# try:\n#     pipe.enable_xformers_memory_efficient_attention()\n# except Exception:\n#     pass\n\n# prompt = \"a high-detail photo of a green pasture under natural sunlight, 35mm, realistic\"\n# negative = \"text, watermark, low quality, blurry\"\n\n# image = pipe(\n#     prompt=prompt,\n#     negative_prompt=negative,\n#     num_inference_steps=25,\n#     guidance_scale=5.0,\n#     height=1024,\n#     width=1024,\n# ).images[0]\n\n# image.save(\"/kaggle/working/sdxl_test.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T06:51:45.244650Z","iopub.execute_input":"2025-12-31T06:51:45.245185Z","iopub.status.idle":"2025-12-31T06:55:09.758537Z","shell.execute_reply.started":"2025-12-31T06:51:45.245158Z","shell.execute_reply":"2025-12-31T06:55:09.757707Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 测试数据增强效果","metadata":{}},{"cell_type":"code","source":"# import os\n# import torch\n# from PIL import Image\n# import matplotlib.pyplot as plt\n# from diffusers import AutoPipelineForImage2Image\n\n# img_path = \"/kaggle/input/csiro-biomass/train/ID21377800.jpg\"\n# model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\n# hf_token = os.environ.get(\"HF_TOKEN\")  # Kaggle Secrets 里配的 HF_TOKEN\n\n# def pad_to_multiple_of_8(pil_img: Image.Image):\n#     w, h = pil_img.size\n#     new_w = ((w + 7) // 8) * 8\n#     new_h = ((h + 7) // 8) * 8\n#     if (new_w, new_h) == (w, h):\n#         return pil_img, (w, h)\n\n#     padded = Image.new(\"RGB\", (new_w, new_h))\n#     padded.paste(pil_img, (0, 0))\n\n#     if new_w > w:\n#         right = pil_img.crop((w - 1, 0, w, h)).resize((new_w - w, h))\n#         padded.paste(right, (w, 0))\n#     if new_h > h:\n#         bottom = padded.crop((0, h - 1, new_w, h)).resize((new_w, new_h - h))\n#         padded.paste(bottom, (0, h))\n\n#     return padded, (w, h)\n\n# # 1) 读图\n# orig = Image.open(img_path).convert(\"RGB\")\n# orig_size = orig.size\n# image, (w0, h0) = pad_to_multiple_of_8(orig)\n\n# # 2) 加载 SDXL img2img pipeline\n# pipe = AutoPipelineForImage2Image.from_pretrained(\n#     model_id,\n#     torch_dtype=torch.float16,\n#     use_safetensors=True,\n#     token=hf_token,\n# ).to(\"cuda\")\n\n# pipe.enable_vae_slicing()\n# try:\n#     pipe.enable_xformers_memory_efficient_attention()\n# except Exception:\n#     pass\n\n# # 3) 关键：提示词写“保持不变”，并把 strength 调低\n# # prompt = (\n# #     \"A realistic photo. Keep the exact same scene, composition, viewpoint, \"\n# #     \"biomass amount and vegetation density as the input image. \"\n# #     \"Only extremely subtle natural variation (tiny lighting/sensor noise).\"\n# # )\n# # negative = (\n# #     \"different scene, different viewpoint, crop, zoom, rotate, \"\n# #     \"add plants, remove plants, change biomass, change vegetation density, \"\n# #     \"change season, change colors, artifacts, cartoon, CGI, painting\"\n# # )\n\n# prompt = (\n#     \"A realistic photo of the same vegetation plot with the same biomass amount and vegetation density.\"\n#     \"Keep the vegetation geometry and coverage unchanged (do not add/remove plants or leaves).\"\n#     \"Introduce noticeable changes in lighting and capture conditions: different sun angle, stronger shadows, different exposure/white balance, slight haze.\"\n#     \"Keep it photorealistic.\"\n# )\n# negative = (\n#     \"different scene, different plot, crop, zoom, rotate, perspective warp,\"\n#     \"add plants, remove plants, change canopy coverage, change vegetation density,\"\n#     \"change leaf shapes, artifacts, cartoon, CGI, painting\"\n# )\n\n# # 越小越“像原图”\n# strength = 0.55\n# guidance_scale = 5\n# steps = 30\n\n# seeds = [11, 22, 33, 44]  # 先出 4 张看看效果\n\n# out_dir = \"/kaggle/working/sdxl_aug\"\n# os.makedirs(out_dir, exist_ok=True)\n\n# outputs = []\n# for seed in seeds:\n#     g = torch.Generator(device=\"cuda\").manual_seed(seed)\n#     out = pipe(\n#         prompt=prompt,\n#         negative_prompt=negative,\n#         image=image,\n#         strength=strength,\n#         guidance_scale=guidance_scale,\n#         num_inference_steps=steps,\n#         generator=g,\n#     ).images[0]\n\n#     # 裁回原尺寸，确保“和原图一样大”\n#     out = out.crop((0, 0, orig_size[0], orig_size[1]))\n#     save_path = os.path.join(out_dir, f\"ID1011485656__sdxl_seed{seed}.jpg\")\n#     out.save(save_path, quality=95)\n#     outputs.append(out)\n\n# print(\"Saved to:\", out_dir)\n\n# # 4) 可视化对比：原图 + 增强图\n# plt.figure(figsize=(16, 8))\n# plt.subplot(2, 3, 1); plt.imshow(orig); plt.title(\"orig\"); plt.axis(\"off\")\n# for i, (seed, im) in enumerate(zip(seeds, outputs), start=2):\n#     plt.subplot(2, 3, i); plt.imshow(im); plt.title(f\"sdxl seed {seed}\"); plt.axis(\"off\")\n# plt.tight_layout()\n# plt.show()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T08:51:34.780994Z","iopub.execute_input":"2025-12-31T08:51:34.781650Z","iopub.status.idle":"2025-12-31T08:57:02.718942Z","shell.execute_reply.started":"2025-12-31T08:51:34.781620Z","shell.execute_reply":"2025-12-31T08:57:02.718099Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 增强数据","metadata":{}},{"cell_type":"code","source":"import gc, torch\n\n# 如果你有 pipe / 大张量，先 del 掉它们\nfor name in [\"pipe\", \"out\", \"orig\", \"g\"]:\n    if name in globals():\n        del globals()[name]\n\ngc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()\n    torch.cuda.ipc_collect()\n\nprint(\"cuda available:\", torch.cuda.is_available())\nif torch.cuda.is_available():\n    print(\"allocated:\", torch.cuda.memory_allocated() / 1024**2, \"MB\")\n    print(\"reserved :\", torch.cuda.memory_reserved() / 1024**2, \"MB\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:03:12.759434Z","iopub.execute_input":"2025-12-31T10:03:12.760129Z","iopub.status.idle":"2025-12-31T10:03:13.118425Z","shell.execute_reply.started":"2025-12-31T10:03:12.760097Z","shell.execute_reply":"2025-12-31T10:03:13.117576Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\nfrom pathlib import Path\nfrom PIL import Image\nimport torch\nfrom diffusers import AutoPipelineForImage2Image\n\ntry:\n    from tqdm.auto import tqdm\nexcept Exception:\n    tqdm = lambda x, **kwargs: x\n\nbase_dir = \"/kaggle/input/csiro-biomass\"\ncsv_in = \"/kaggle/input/csiro-biomass/train.csv\"\n\nout_images_dir = \"/kaggle/working/csiro-biomass-images\"\ncsv_out = \"/kaggle/working/train_aug_long.csv\"\nos.makedirs(out_images_dir, exist_ok=True)\n\nmodel_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\nhf_token = os.environ.get(\"HF_TOKEN\")\n\nprompt = (\n    \"A realistic photo of the same vegetation plot with the same biomass amount and vegetation density. \"\n    \"Keep the vegetation geometry and coverage unchanged (do not add/remove plants or leaves). \"\n    \"Introduce noticeable changes in lighting and capture conditions: different sun angle, stronger shadows, \"\n    \"different exposure/white balance, slight haze. Keep it photorealistic.\"\n)\nnegative = (\n    \"different scene, different plot, crop, zoom, rotate, perspective warp, \"\n    \"add plants, remove plants, change canopy coverage, change vegetation density, \"\n    \"change leaf shapes, artifacts, cartoon, CGI, painting\"\n)\n\nstrength = 0.35\nguidance_scale = 4.5\nsteps = 30\nseeds = [11, 22]  # 每张原图生成 2 张\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"device:\", device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:03:17.527841Z","iopub.execute_input":"2025-12-31T10:03:17.528406Z","iopub.status.idle":"2025-12-31T10:03:18.713898Z","shell.execute_reply.started":"2025-12-31T10:03:17.528376Z","shell.execute_reply":"2025-12-31T10:03:18.713302Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pipe = AutoPipelineForImage2Image.from_pretrained(\n    model_id,\n    torch_dtype=torch.float16 if device == \"cuda\" else torch.float32,\n    use_safetensors=True,\n    token=hf_token,\n).to(device)\n\npipe.enable_vae_slicing()\n\n# 关键：VAE 用 fp32 解码更稳，常见可直接解决“全黑图”\nif hasattr(pipe, \"upcast_vae\"):\n    pipe.upcast_vae()\nelse:\n    pipe.vae.to(dtype=torch.float32)\n\n# 默认不启用 xformers（很多“能跑但黑图”就是它引入的数值问题）\n# 如果你确认启用没问题再打开：\n# try:\n#     pipe.enable_xformers_memory_efficient_attention()\n# except Exception:\n#     pass\n\npipe.set_progress_bar_config(disable=True)\nprint(\"pipeline loaded; unet dtype:\", pipe.unet.dtype, \"vae dtype:\", pipe.vae.dtype)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:03:19.892754Z","iopub.execute_input":"2025-12-31T10:03:19.893946Z","iopub.status.idle":"2025-12-31T10:05:47.666506Z","shell.execute_reply.started":"2025-12-31T10:03:19.893910Z","shell.execute_reply":"2025-12-31T10:05:47.665828Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"USECOLS = [\n    \"sample_id\", \"image_path\", \"target_name\", \"target\",\n    \"State\", \"Species\", \"Height_Ave_cm\", \"Sampling_Date\",\n]\ndf = pd.read_csv(csv_in, usecols=USECOLS)\n\ndef image_id_from_sample_id(s):\n    return str(s).split(\"_\")[0]\n\ndef replace_prefix_sample_id(sample_id, new_image_id):\n    s = str(sample_id)\n    if \"_\" in s:\n        _, rest = s.split(\"_\", 1)\n        return f\"{new_image_id}_{rest}\"\n    return new_image_id\n\ndf[\"image_id\"] = df[\"sample_id\"].map(image_id_from_sample_id)\n\n# 每个 image_id 取一个相对路径来读图（同一张图在长表里会出现多行）\nid_to_relpath = df.groupby(\"image_id\")[\"image_path\"].first().to_dict()\n\nprint(\"rows:\", len(df), \"unique images:\", len(id_to_relpath))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:05:47.667775Z","iopub.execute_input":"2025-12-31T10:05:47.668050Z","iopub.status.idle":"2025-12-31T10:05:47.715547Z","shell.execute_reply.started":"2025-12-31T10:05:47.667995Z","shell.execute_reply":"2025-12-31T10:05:47.714931Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"new_rows = df[USECOLS].to_dict(orient=\"records\")  # 先保留原始所有行\n\nfailed = []\nfirst_debug_printed = False\n\nfor image_id, rel_path in tqdm(id_to_relpath.items(), desc=\"augment\"):\n    img_path = str(Path(base_dir) / str(rel_path))\n    try:\n        orig = Image.open(img_path).convert(\"RGB\")\n    except Exception as e:\n        failed.append((image_id, img_path, str(e)))\n        continue\n\n    # 该 image_id 对应的所有标签行（长格式多行）\n    sub = df.loc[df[\"image_id\"] == image_id, USECOLS].copy()\n\n    for k, seed in enumerate(seeds, start=1):\n        new_image_id = f\"{image_id}__sdxl{k}\"\n        save_path = os.path.join(out_images_dir, f\"{new_image_id}.jpg\")\n\n        g = torch.Generator(device=device).manual_seed(int(seed))\n\n        out_img = pipe(\n            prompt=prompt,\n            negative_prompt=negative,\n            image=orig,  # 你说原始尺寸可被 8 整除：不做 pad/crop\n            strength=strength,\n            guidance_scale=guidance_scale,\n            num_inference_steps=steps,\n            generator=g,\n        ).images[0]\n\n        # 只打印一次，用来确认不是黑图（可保留也可删掉）\n        if not first_debug_printed:\n            import numpy as np\n            arr = np.asarray(out_img)\n            print(\"debug out min/max/mean:\", arr.min(), arr.max(), float(arr.mean()))\n            first_debug_printed = True\n\n        out_img.save(save_path, quality=95)\n\n        # 复制该图所有 target 行到新图（保持长格式）\n        sub_k = sub.copy()\n        sub_k[\"sample_id\"] = sub_k[\"sample_id\"].map(lambda s: replace_prefix_sample_id(s, new_image_id))\n        sub_k[\"image_path\"] = save_path\n        new_rows.extend(sub_k.to_dict(orient=\"records\"))\n\n    if device == \"cuda\":\n        torch.cuda.empty_cache()\n\ndf_aug_long = pd.DataFrame(new_rows, columns=USECOLS)\ndf_aug_long.to_csv(csv_out, index=False)\n\nprint(\"saved images to:\", out_images_dir)\nprint(\"saved csv to:\", csv_out)\nprint(\"rows:\", len(df), \"->\", len(df_aug_long))\nprint(\"failed:\", len(failed))\nif failed:\n    print(\"first failed:\", failed[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-31T10:05:47.716348Z","iopub.execute_input":"2025-12-31T10:05:47.716627Z","execution_failed":"2025-12-31T10:12:03.677Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}