diff --git a/Makefile b/Makefile index 2df576a7..a6787a88 100644 --- a/Makefile +++ b/Makefile @@ -82,11 +82,11 @@ constraints: deepclean # Check lint with black. black: - $(PIPRUN) python -m black --check --diff . --extend-exclude test/scripts --extend-exclude git_ignore_folder -l 120 + $(PIPRUN) python -m black --check --diff . --extend-exclude "(test/scripts|test/notebook/testfiles|git_ignore_folder)" -l 120 # Check lint with isort. isort: - $(PIPRUN) python -m isort --check . -s git_ignore_folder -s test/scripts + $(PIPRUN) python -m isort --check . -s git_ignore_folder -s test/scripts -s test/notebook/testfiles # Check lint with mypy. # First deal with the core folder, and then gradually increase the scope of detection, @@ -119,11 +119,11 @@ pre-commit: # Auto lint with black. auto-black: - $(PIPRUN) python -m black . --extend-exclude test/scripts --extend-exclude git_ignore_folder --extend-exclude .venv -l 120 + $(PIPRUN) python -m black . --extend-exclude "(test/scripts|test/notebook/testfiles|git_ignore_folder|.venv)" -l 120 # Auto lint with isort. auto-isort: - $(PIPRUN) python -m isort . -s git_ignore_folder -s test/scripts -s .venv + $(PIPRUN) python -m isort . -s git_ignore_folder -s test/scripts -s test/notebook/testfiles -s .venv # Auto lint with toml-sort. auto-toml-sort: diff --git a/rdagent/app/data_science/conf.py b/rdagent/app/data_science/conf.py index e24942df..49591074 100644 --- a/rdagent/app/data_science/conf.py +++ b/rdagent/app/data_science/conf.py @@ -43,6 +43,9 @@ class DataScienceBasePropSetting(KaggleBasePropSetting): ### specific feature + ### notebook integration + enable_notebook_conversion: bool = False + #### enable specification spec_enabled: bool = True diff --git a/rdagent/components/coder/data_science/pipeline/__init__.py b/rdagent/components/coder/data_science/pipeline/__init__.py index 38f99efd..8a5f27d2 100644 --- a/rdagent/components/coder/data_science/pipeline/__init__.py +++ b/rdagent/components/coder/data_science/pipeline/__init__.py @@ -83,7 +83,10 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): package_info=target_task.package_info, enable_model_dump=DS_RD_SETTING.enable_model_dump, enable_debug_mode=DS_RD_SETTING.sample_data_by_LLM, - spec=T("scenarios.data_science.share:component_spec.Pipeline").r(metric_name=self.scen.metric_name), + spec=T("scenarios.data_science.share:component_spec.Pipeline").r( + metric_name=self.scen.metric_name, + enable_notebook_conversion=DS_RD_SETTING.enable_notebook_conversion, + ), ) user_prompt = T(".prompts:pipeline_coder.user").r( competition_info=competition_info, diff --git a/rdagent/components/coder/data_science/pipeline/eval.py b/rdagent/components/coder/data_science/pipeline/eval.py index 3d8cbf4e..c7b92dc5 100644 --- a/rdagent/components/coder/data_science/pipeline/eval.py +++ b/rdagent/components/coder/data_science/pipeline/eval.py @@ -16,6 +16,7 @@ from rdagent.components.coder.CoSTEER.knowledge_management import ( CoSTEERQueriedKnowledgeV2, ) from rdagent.components.coder.data_science.conf import get_clear_ws_cmd, get_ds_env +from rdagent.components.coder.data_science.share.notebook import NotebookConverter from rdagent.components.coder.data_science.utils import remove_eda_part from rdagent.core.experiment import FBWorkspace, Task from rdagent.scenarios.data_science.test_eval import get_test_eval @@ -70,6 +71,24 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): env=env, entry=f"strace -e trace=file -f -o trace.log python -m coverage run main.py" ) + nb_conversion_ret_code = 0 + nb_conversion_check_text = "" + if DS_RD_SETTING.enable_notebook_conversion: + notebook_converter = NotebookConverter() + code = implementation.file_dict["main.py"] + error_msg = notebook_converter.validate_code_format(code) + if error_msg is not None: + nb_conversion_check_text = error_msg + nb_conversion_ret_code = 1 + else: + notebook_converter.convert( + task=target_task, + code=code, + stdout=result.stdout, + outfile=implementation.workspace_path / "main.ipynb", + use_debug_flag=DS_RD_SETTING.sample_data_by_LLM, + ) + sample_submission_check = True test_eval = get_test_eval() if (sample_submission_file_name := test_eval.get_sample_submission_name(self.scen.competition)) is not None: @@ -173,7 +192,10 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): scenario=self.scen.get_scenario_all_desc(eda_output=eda_output), task_desc=target_task.get_task_information(), stdout=stdout.strip(), - spec=T("scenarios.data_science.share:component_spec.Pipeline").r(metric_name=self.scen.metric_name), + spec=T("scenarios.data_science.share:component_spec.Pipeline").r( + metric_name=self.scen.metric_name, + enable_notebook_conversion=DS_RD_SETTING.enable_notebook_conversion, + ), code=implementation.file_dict["main.py"], ) wfb = build_cls_from_json_with_retry( @@ -193,4 +215,7 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): wfb.return_checking += ( "\nSample submission file check failed. Code should not open the sample submission file." ) + if nb_conversion_ret_code != 0 and wfb.final_decision is True: + wfb.final_decision = False + wfb.return_checking += "\n" + nb_conversion_check_text return wfb diff --git a/rdagent/components/coder/data_science/share/notebook.py b/rdagent/components/coder/data_science/share/notebook.py new file mode 100644 index 00000000..09a909fb --- /dev/null +++ b/rdagent/components/coder/data_science/share/notebook.py @@ -0,0 +1,135 @@ +""" +Handles conversion from a Python file to a Jupyter notebook. +""" + +import argparse +from typing import Optional + +import nbformat + +from rdagent.components.coder.data_science.share.util import ( + extract_first_section_name_from_code, + extract_function_body, + split_code_and_output_into_sections, +) +from rdagent.core.experiment import Task +from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_utils import APIBackend +from rdagent.utils.agent.ret import MarkdownAgentOut +from rdagent.utils.agent.tpl import T + + +class NotebookConverter: + """ + Builder responsible for writing a Jupyter notebook for a workspace. + """ + + def validate_code_format(self, code: str) -> str | None: + """ + Returns None if the code format is valid, otherwise returns an error message. + """ + main_function_body = extract_function_body(code, "main") + if not main_function_body: + return "[Error] No main function found in the code. Please ensure that the main function is defined and contains the necessary print statements to divide sections." + + found_section_name = extract_first_section_name_from_code(main_function_body) + if not found_section_name: + return "[Error] No sections found in the code. Expected to see 'print(\"Section:
\")' as section dividers. Also make sure that they are actually run and not just comments." + + return None + + def convert( + self, + task: Optional[Task], + code: str, + stdout: str, + outfile: Optional[str] = None, + use_debug_flag: bool = False, + ) -> str: + """ + Build a notebook based on the current progression. + """ + # Handle argparse in the code to ensure it works in a notebook environment + should_handle_argparse = "argparse" in code + sections = split_code_and_output_into_sections(code=code, stdout=stdout) + notebook = nbformat.v4.new_notebook() + + # Use LLM to generate an intro cell for the notebook + if task: + system_prompt = T(".prompts:notebookconverter.system").r() + user_prompt = T(".prompts:notebookconverter.user").r( + plan=task.get_task_information(), + code=code, + ) + resp = APIBackend().build_messages_and_create_chat_completion( + user_prompt=user_prompt, system_prompt=system_prompt + ) + intro_content = MarkdownAgentOut.extract_output(resp) + notebook.cells.append(nbformat.v4.new_markdown_cell(intro_content)) + + if should_handle_argparse: + # Remove extra `import sys` since it will be added for argparse handling + if "import sys\n" in sections[0]["code"]: + sections[0]["code"] = sections[0]["code"].replace("import sys\n", "") + + # Add sys.argv modification for argparse handling + sections[0]["code"] = ( + "\n".join( + [ + "import sys", + "# hack to allow argparse to work in notebook", + ('sys.argv = ["main.py", "--debug"]' if use_debug_flag else 'sys.argv = ["main.py"]'), + ] + ) + + "\n\n" + + sections[0]["code"].lstrip() + ) + + for section in sections: + # Create a markdown cell for the section name and comments + markdown_content = "" + if section["name"]: + markdown_content += f"## {section['name']}\n" + if section["comments"]: + markdown_content += f"{section['comments']}\n" + if markdown_content: + notebook.cells.append(nbformat.v4.new_markdown_cell(markdown_content)) + + # Create a code cell for the section code and output + if section["code"]: + cell = nbformat.v4.new_code_cell(section["code"]) + if section["output"]: + # For simplicity, treat all output as coming from stdout + # TODO: support Jupyter kernel execution and handle outputs appropriately here + cell.outputs = [nbformat.v4.new_output("stream", name="stdout", text=section["output"])] + notebook.cells.append(cell) + + # Save the notebook or return it as a string + if outfile: + with open((outfile), "w", encoding="utf-8") as f: + nbformat.write(notebook, f) + logger.info(f"Notebook written to {outfile}") + + return nbformat.writes(notebook) + + +if __name__ == "__main__": + converter = NotebookConverter() + parser = argparse.ArgumentParser(description="Convert Python code to Jupyter notebook.") + parser.add_argument("inputfile", type=str, help="Path to the input Python file.") + parser.add_argument("outfile", type=str, help="Path to the output Notebook file.") + parser.add_argument( + "--stdout", + type=str, + default="", + help="Standard output from the code execution.", + ) + parser.add_argument("--debug", action="store_true", help="Use debug flag to modify sys.argv.") + args = parser.parse_args() + converter.convert( + task=None, + code=open(args.inputfile, "r").read(), + stdout=args.stdout, + outfile=args.outfile, + use_debug_flag=False, + ) diff --git a/rdagent/components/coder/data_science/share/prompts.yaml b/rdagent/components/coder/data_science/share/prompts.yaml index 3bf09147..3657c160 100644 --- a/rdagent/components/coder/data_science/share/prompts.yaml +++ b/rdagent/components/coder/data_science/share/prompts.yaml @@ -89,3 +89,25 @@ docdev: ``` {% endfor %} +notebookconverter: + system: |- + {% include "scenarios.data_science.share:scen.role" %} Your task is to provide a summary for a data science solution. + + You will be given: + - The original implementation plan for the script. + - A Python script that contains code and output. + + Your task is to generate markdown content that includes a title and a short paragraph summarizing the technique in model training, the type of model produced and any other noteworthy details in the solution. + + The return content should be like the format below(Please note that "````" is used to avoid confliction of "```" in markdown file) + ````markdown + # + + ```` + + user: |- + --------------- The implementation plan --------------- + {{plan}} + + --------------- The Python script content --------------- + {{code}} diff --git a/rdagent/components/coder/data_science/share/util.py b/rdagent/components/coder/data_science/share/util.py new file mode 100644 index 00000000..5f6f9a61 --- /dev/null +++ b/rdagent/components/coder/data_science/share/util.py @@ -0,0 +1,365 @@ +import ast +import io +import re +import tokenize +from itertools import zip_longest +from typing import List, Optional, Set, Tuple, TypedDict + + +class CodeSection(TypedDict): + """ + Represents a section of the original Python source code, to be converted to a notebook cell. + """ + + name: Optional[str] + code: Optional[str] + comments: Optional[str] + output: Optional[str] + + +def extract_function_body(source_code: str, function_name: str) -> Optional[str]: + """ + Extracts the body of a function from the source code. + Returns None if the function is not found. + + Assumption: The function is multiline and defined at the top level. + """ + tree = ast.parse(source_code) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == function_name: + lines = source_code.splitlines() + start = node.body[0].lineno + end = node.body[-1].end_lineno + body_lines = lines[start - 1 : end] + indent_level = len(body_lines[0]) - len(body_lines[0].lstrip()) + return "\n".join(line[indent_level:] for line in body_lines) + return None + + +def split_sections( + text: str, section_header_regex: str, known_sections: Optional[list[str]] = None +) -> tuple[Optional[str], list[str], list[str]]: + """ + Split text into sections based on the section headers. + """ + sections = [] + section_names = [] + current_section = [] + next_section_name_index = 0 + for line in text.splitlines(): + match = re.match(section_header_regex, line) + extracted_section_name = match.group(1).strip() if match else None + if extracted_section_name and ( + not known_sections + or ( + next_section_name_index < len(known_sections) + and extracted_section_name == known_sections[next_section_name_index] + ) + ): + if current_section: + sections.append("\n".join(current_section)) + current_section = [] + current_section.append(line) + section_names.append(extracted_section_name) + next_section_name_index += 1 + else: + current_section.append(line) + if current_section: + sections.append("\n".join(current_section)) + + # If the first section does not match the header regex, treat it as a header section. + header_section = None + if sections and not re.search(section_header_regex, sections[0]): + header_section = sections[0] + sections = sections[1:] + + return header_section, sections, section_names + + +def split_code_sections(source_code: str) -> tuple[Optional[str], list[str]]: + """ + Split code into sections based on the section headers. + """ + return split_sections(source_code, r'^print\(["\']Section: (.+)["\']\)') + + +def split_output_sections(stdout: str, known_sections: list[str]) -> tuple[Optional[str], list[str]]: + """ + Split output into sections based on the section headers. + """ + header_section, sections, _ = split_sections(stdout, r"^Section: (.+)", known_sections=known_sections) + return header_section, sections + + +def extract_comment_under_first_print(source_code) -> tuple[Optional[str], str]: + """ + Extract comments from the source code after the first print statement. + """ + lines = source_code.splitlines() + lines_to_remove = set() + all_comments = [] + + parsed = ast.parse(source_code) + # Find the first print statement only + first_print_lineno = None + for node in ast.walk(parsed): + if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + if getattr(node.value.func, "id", None) == "print": + first_print_lineno = node.lineno + break + + if first_print_lineno is None: + # No print statement found, return empty comments and original code + return None, source_code + + for i in range(first_print_lineno, len(lines)): + stripped = lines[i].strip() + if stripped.startswith("#"): + comment_text = stripped.lstrip("# ").strip() + all_comments.append(comment_text) + lines_to_remove.add(i) + elif stripped == "": + continue + elif i > first_print_lineno: + break # stop after hitting actual code line + + cleaned_lines = [line for idx, line in enumerate(lines) if idx not in lines_to_remove] + cleaned_code = "\n".join(cleaned_lines) + comments_str = "\n".join(all_comments) if all_comments else None + + return comments_str, cleaned_code + + +def extract_first_section_name_from_code(source_code): + """ + Extract the first section name from the source code. + """ + parsed = ast.parse(source_code) + for node in ast.walk(parsed): + if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + call = node.value + if getattr(call.func, "id", None) == "print" and call.args: + arg0 = call.args[0] + if isinstance(arg0, ast.Constant) and isinstance(arg0.value, str): + # Match "Section: ..." pattern + m = re.match(r"Section:\s*(.+)", arg0.value) + if m: + return m.group(1).strip() + return None + + +def extract_first_section_name_from_output(stdout: str) -> Optional[str]: + """ + Extract the first section name from the output string. + """ + match = re.search(r"Section:\s*(.+)", stdout) + if match: + return match.group(1).strip() + return None + + +def is_function_called(source_code: str, func_name: str) -> bool: + """ + Returns True if the function named `func_name` is called in `source_code`. + """ + tree = ast.parse(source_code) + for node in ast.walk(tree): + if isinstance(node, ast.Call): + # For simple function calls like func() + if isinstance(node.func, ast.Name) and node.func.id == func_name: + return True + + # For calls like module.func() + elif isinstance(node.func, ast.Attribute) and node.func.attr == func_name: + return True + return False + + +def remove_function(source_code: str, function_name: str) -> str: + """ + Remove a function definition from the source code. + """ + tree = ast.parse(source_code) + lines = source_code.splitlines() + + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == function_name: + start_lineno = node.lineno - 1 + end_lineno = node.end_lineno + return "\n".join(lines[:start_lineno] + lines[end_lineno:]) + + return source_code + + +def remove_main_block(source_code: str) -> str: + """ + Remove the if __name__ == "__main__": block from the source code. + """ + tree = ast.parse(source_code) + lines = source_code.splitlines() + + # Find the main block and note its line numbers + for node in tree.body: + if isinstance(node, ast.If): + test = node.test + if ( + isinstance(test, ast.Compare) + and isinstance(test.left, ast.Name) + and test.left.id == "__name__" + and len(test.ops) == 1 + and isinstance(test.ops[0], ast.Eq) + and len(test.comparators) == 1 + and isinstance(test.comparators[0], ast.Constant) + and test.comparators[0].value == "__main__" + ): + + # Remove lines corresponding to this block + start_lineno = node.lineno - 1 + end_lineno = node.end_lineno + return "\n".join(lines[:start_lineno] + lines[end_lineno:]) + + return source_code + + +def extract_top_level_functions_with_decorators_and_comments( + code: str, +) -> List[Tuple[str, str]]: + """ + Returns list of (function_name, source_segment) for top-level functions (excluding "main"), + including decorators and contiguous preceding comments. + """ + # Parse AST to get function nodes + tree = ast.parse(code) + lines = code.splitlines(keepends=True) + + # Precompute which line numbers have comment tokens + comment_lines: Set[int] = set() + lines = code.splitlines(keepends=True) # preserve exact line content for prefix checks + + tokgen = tokenize.generate_tokens(io.StringIO(code).readline) # yields (type, string, start, end, line) + for tok_type, _, (srow, scol), _, _ in tokgen: + if tok_type == tokenize.COMMENT: + # everything before the comment on that line must be whitespace + prefix = lines[srow - 1][:scol] + if prefix.strip() == "": + comment_lines.add(srow) + + functions = [] + + for node in tree.body: # only top-level + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + if node.name == "main": + continue + + # Determine the starting line: earliest decorator if present, else the def/async line + if node.decorator_list: + start_lineno = min(d.lineno for d in node.decorator_list) + else: + start_lineno = node.lineno + + # Extend upward to include contiguous comment lines (no intervening non-blank/non-comment) + span_start = start_lineno + curr = span_start - 1 # check line above; lines are 1-based + while curr > 0: + line_text = lines[curr - 1] + if curr in comment_lines: + span_start = curr + curr -= 1 + continue + if line_text.strip() == "": + # blank line: include it and keep scanning upward + span_start = curr + curr -= 1 + continue + break # encountered code or something else; stop + + # Determine end line of the function definition including its body + # Prefer end_lineno if available (Python 3.8+) + if hasattr(node, "end_lineno") and node.end_lineno is not None: + span_end = node.end_lineno + else: + # Fallback: get last lineno from the deepest child in body + def _max_lineno(n): + max_ln = getattr(n, "lineno", 0) + for child in ast.iter_child_nodes(n): + ln = _max_lineno(child) + if ln > max_ln: + max_ln = ln + return max_ln + + span_end = _max_lineno(node) + + # Slice the original source lines + segment = "".join(lines[span_start - 1 : span_end]) + functions.append((node.name, segment)) + + return functions + + +def split_code_and_output_into_sections(code: str, stdout: str) -> list[CodeSection]: + """ + Converts a Python script and its output into a list of CodeSections. + Pre-condition: The code in the main() function contains print statements that indicate section names, e.g., `print("Section:
")`. + """ + # This will hold all top-level code and by default all function definitions. + # Functions will later be moved to more relevant sections if needed. + # The first step is to remove both the if __name__ == "__main__": block and the main function + top_level_code = remove_main_block(remove_function(code, "main")) + + main_function_body = extract_function_body(code, "main") + functions = extract_top_level_functions_with_decorators_and_comments(top_level_code) + + # Split the main function body into sections based on print("Section:
") code + main_fn_top_level_section, main_fn_sections, known_section_names = ( + split_code_sections(main_function_body) if main_function_body else (None, [], []) + ) + + # Split the output into sections based on "Section: " headers + output_top_level_section, output_sections = split_output_sections(stdout, known_section_names) + + # Merge code and outputs into code sections + result_sections: list[CodeSection] = [] + for output_section, code_section in zip_longest(output_sections, main_fn_sections): + name = None + if code_section is not None: + # If code section is available, extract the section name from it + name = extract_first_section_name_from_code(code_section) + elif output_section: + # If only output section is available, extract the section name from it + name = extract_first_section_name_from_output(output_section) + comments, cleaned_code = ( + extract_comment_under_first_print(code_section) if code_section is not None else (None, None) + ) + # Strip whitespaces for the cell + if cleaned_code is not None: + cleaned_code = cleaned_code.strip() + result_sections.append(CodeSection(name=name, code=cleaned_code, comments=comments, output=output_section)) + + # Small optimization: move function definitions to the sections where they are first called + # TODO: this doesn't handle nested function references, e.g., fn A calls fn B which calls fn C + # currently will not move C to the section where A is called + for name, segment in functions: + for section in result_sections: + if section["code"] and is_function_called(section["code"], name): + section["code"] = segment.strip() + "\n\n" + section["code"].lstrip() + top_level_code = top_level_code.replace(segment, "") + break + + # Inject the top-level code at the beginning of the sections + top_level_code = ( + top_level_code.rstrip() + "\n\n" + main_fn_top_level_section.lstrip() + if main_fn_top_level_section + else top_level_code + ) + result_sections.insert( + 0, + CodeSection( + name=None, + code=top_level_code, + comments=None, + output=output_top_level_section, + ), + ) + + return result_sections diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/draft/draft.py b/rdagent/scenarios/data_science/proposal/exp_gen/draft/draft.py index 5e51faf1..1a60f1e7 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/draft/draft.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/draft/draft.py @@ -106,7 +106,9 @@ class DSDraftExpGen(ExpGen): if DS_RD_SETTING.spec_enabled: spec = last_successful_exp.experiment_workspace.file_dict[spec_file] if spec_file else None else: - spec = T(f"scenarios.data_science.share:component_spec.{component}").r() + spec = T(f"scenarios.data_science.share:component_spec.{component}").r( + enable_notebook_conversion=DS_RD_SETTING.enable_notebook_conversion, + ) resp_dict = self._init_task_gen( targets=component, scenario_desc=scenario_desc, diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py index bea8b34d..55a36664 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py @@ -388,7 +388,9 @@ class DSProposalV1ExpGen(ExpGen): if DS_RD_SETTING.spec_enabled: task_spec = sota_exp.experiment_workspace.file_dict[component_info["spec_file"]] else: - task_spec = T(f"scenarios.data_science.share:component_spec.{component}").r() + task_spec = T(f"scenarios.data_science.share:component_spec.{component}").r( + enable_notebook_conversion=DS_RD_SETTING.enable_notebook_conversion, + ) system_prompt = T(".prompts:direct_exp_gen.system").r( targets=component_info["target_name"], component=component, diff --git a/rdagent/scenarios/data_science/share.yaml b/rdagent/scenarios/data_science/share.yaml index 2bda21e0..681d7f9f 100644 --- a/rdagent/scenarios/data_science/share.yaml +++ b/rdagent/scenarios/data_science/share.yaml @@ -268,7 +268,44 @@ component_spec: 1. Program Execution: - The workflow will be executed by running `python main.py` with no command-line arguments. Ensure that `main.py` does not require or expect any parameters. - The working directory will only contain `main.py`. Any additional files required for execution must be downloaded or generated by `main.py` itself. - + {% if enable_notebook_conversion %} + - Code should be modular and organized into functions, with a clear main() function that orchestrates the workflow. + - Inside the main() function, divide the code into sequential sections. + - Each section must follow this exact pattern: + - A print statement announcing the section, e.g. print("Section: ") + - 1–2 lines of comments explaining the purpose of the section + - The block of code for that section + - Example: + ```python + + + def main(): + print("Section: Data Loading") + # Load dataset from CSV into a DataFrame + # Handle missing values + + + print("Section: Feature Engineering") + # Generate new features from raw data + # Normalize and encode categorical variables + + + + + + + if __name__ == "__main__": + main() + ``` + - Section headers must always be print("Section: ") at the top level of main(). + - Do not put them inside if/else blocks. + - Do not put them inside helper functions. + - Do not put them outside of the main() function. + - Do not add comments or dividers before the print statement — comments must come after the print line. + - Section names should be concise and descriptive (e.g., "Data Loading", "Model Training"). + - Do not call return or exit in the middle of the main() function. Raise an exception if you need to stop execution early. + {% endif %} + 2. File Handling: - Handle file encoding and delimiters appropriately. - Combine or process multiple files if necessary. diff --git a/requirements.txt b/requirements.txt index 21bc42eb..a6b0f6e5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -52,7 +52,7 @@ networkx # kaggle crawler selenium kaggle -nbformat +nbformat # also used for notebook conversion # tool setuptools-scm diff --git a/test/notebook/test_notebook_converter.py b/test/notebook/test_notebook_converter.py new file mode 100644 index 00000000..811ed239 --- /dev/null +++ b/test/notebook/test_notebook_converter.py @@ -0,0 +1,186 @@ +import json +import os +import unittest + +from rdagent.components.coder.data_science.share.notebook import NotebookConverter + +test_files_dir = os.path.join(os.path.dirname(__file__), "testfiles") + + +def normalize_nb_json_for_comparison(nb_json_str): + nb_json = json.loads(nb_json_str) + for cell in nb_json["cells"]: + if "id" in cell: + cell.pop("id", None) + return json.dumps(nb_json, indent=4) + + +class TestNotebookConverter(unittest.TestCase): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.converter = NotebookConverter() + self.maxDiff = None + + def test_validation_pass(self): + with open(os.path.join(test_files_dir, "main.py"), "r") as f: + code = f.read() + result = self.converter.validate_code_format(code) + self.assertIsNone(result, "Code format should be valid") + + def test_validation_missing_main_fn(self): + with open(os.path.join(test_files_dir, "main_missing_main_fn.py"), "r") as f: + code = f.read() + result = self.converter.validate_code_format(code) + self.assertEqual( + result, + "[Error] No main function found in the code. Please ensure that the main function is defined and contains the necessary print statements to divide sections.", + ) + + def test_validation_missing_sections(self): + with open(os.path.join(test_files_dir, "main_missing_sections.py"), "r") as f: + code = f.read() + result = self.converter.validate_code_format(code) + self.assertEqual( + result, + "[Error] No sections found in the code. Expected to see 'print(\"Section:
\")' as section dividers. Also make sure that they are actually run and not just comments.", + ) + + def test_argparse_happy_path(self): + code = """import argparse +parser = argparse.ArgumentParser(description='Test script') +parser.add_argument('--debug', action='store_true', help='Enable debug mode') +args = parser.parse_args() + +def main(): + print(args.debug) + print("Section: Data Loading") + # Load dataset from CSV into a DataFrame + load_data() + +if __name__ == "__main__": + main()""" + notebookJson = json.loads( + self.converter.convert( + task=None, + code=code, + stdout="", + use_debug_flag=True, + ) + ) + self.assertEqual( + "".join(notebookJson["cells"][0]["source"]), + """import sys +# hack to allow argparse to work in notebook +sys.argv = ["main.py", "--debug"] + +import argparse +parser = argparse.ArgumentParser(description='Test script') +parser.add_argument('--debug', action='store_true', help='Enable debug mode') +args = parser.parse_args() + +print(args.debug)""", + ) + + self.assertEqual( + "".join(notebookJson["cells"][1]["source"]), + """## Data Loading +Load dataset from CSV into a DataFrame +""", + ) + self.assertEqual( + "".join(notebookJson["cells"][2]["source"]), + """print("Section: Data Loading") +load_data()""", + ) + + def test_argparse_with_dupe_sys(self): + code = """import argparse +import sys +parser = argparse.ArgumentParser(description='Test script') +parser.add_argument('--debug', action='store_true', help='Enable debug mode') +args = parser.parse_args() + +print(sys) + +def main(): + print(args.debug) + print("Section: Data Loading") + # Load dataset from CSV into a DataFrame + load_data() + +if __name__ == "__main__": + main()""" + notebookJson = json.loads( + self.converter.convert( + task=None, + code=code, + stdout="", + use_debug_flag=True, + ) + ) + self.assertEqual( + "".join(notebookJson["cells"][0]["source"]), + """import sys +# hack to allow argparse to work in notebook +sys.argv = ["main.py", "--debug"] + +import argparse +parser = argparse.ArgumentParser(description='Test script') +parser.add_argument('--debug', action='store_true', help='Enable debug mode') +args = parser.parse_args() + +print(sys) + +print(args.debug)""", + ) + + self.assertEqual( + "".join(notebookJson["cells"][1]["source"]), + """## Data Loading +Load dataset from CSV into a DataFrame +""", + ) + self.assertEqual( + "".join(notebookJson["cells"][2]["source"]), + """print("Section: Data Loading") +load_data()""", + ) + + def test_convert(self): + with open(os.path.join(test_files_dir, "main.py"), "r") as f: + code = f.read() + notebookJson = self.converter.convert( + task=None, + code=code, + stdout="", + # outfile=os.path.join(test_files_dir, "main.ipynb"), # Uncomment this to save to the file + ) + with open(os.path.join(test_files_dir, "main.ipynb"), "r") as f: + expected_notebook = f.read() + self.assertEqual( + normalize_nb_json_for_comparison(notebookJson), + normalize_nb_json_for_comparison(expected_notebook), + "Converted notebook should match expected output", + ) + + def test_convert_2(self): + with open(os.path.join(test_files_dir, "main2.py"), "r") as f: + code = f.read() + notebookJson = self.converter.convert( + task=None, + code=code, + stdout="", + # outfile=os.path.join(test_files_dir, "main2.ipynb"), # Uncomment this to save to the file + ) + with open(os.path.join(test_files_dir, "main2.ipynb"), "r") as f: + expected_notebook = f.read() + self.assertEqual( + normalize_nb_json_for_comparison(notebookJson), + normalize_nb_json_for_comparison(expected_notebook), + "Converted notebook should match expected output", + ) + + +if __name__ == "__main__": + unittest.main() + # pytest test/notebook/test_notebook_converter.py diff --git a/test/notebook/test_util.py b/test/notebook/test_util.py new file mode 100644 index 00000000..5fee8971 --- /dev/null +++ b/test/notebook/test_util.py @@ -0,0 +1,1922 @@ +import os +import unittest + +from rdagent.components.coder.data_science.share.util import ( + extract_comment_under_first_print, + extract_first_section_name_from_code, + extract_first_section_name_from_output, + extract_function_body, + extract_top_level_functions_with_decorators_and_comments, + is_function_called, + remove_function, + remove_main_block, + split_code_and_output_into_sections, + split_code_sections, + split_output_sections, +) + +test_files_dir = os.path.join(os.path.dirname(__file__), "testfiles") + + +class TestExtractFunctionBody(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "def main():", + " print('Section: Data Loading')", + " # Load data", + " data = load_data()", + "", + ] + ) + extracted = extract_function_body(code, "main") + expected = S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + ] + ) + self.assertEqual(extracted, expected) + + def test_happy_path_complex(self): + code = S( + [ + "import pandas as pd", + "", + "print('main()')", + "", + "def foo():", + " print('Section: Foo')", + "", + "def mainfunc():", + " print('Section: Data Loading 2')", + " # Load data 2", + " data2 = load_data()", + "", + "def main():", + " print('Section: Data Loading')", + " # Load data", + " data = load_data()", + "", + "def bar():", + " print('Section: Foo')", + "", + "main()", + ] + ) + extracted = extract_function_body(code, "main") + expected = S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + ] + ) + self.assertEqual(extracted, expected) + + def test_empty(self): + extracted = extract_function_body("", "main") + expected = None + self.assertEqual(extracted, expected) + + def test_missing_func(self): + code = S( + [ + "def foo():", + " print('Section: Data Loading')", + " # Load data", + " data = load_data()", + "", + ] + ) + extracted = extract_function_body(code, "main") + expected = None + self.assertEqual(extracted, expected) + + +class TestSplitCodeSections(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "# This is the main function", + "setup_workspace()", + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + 'print("Section: Data Processing")', + "# Process data", + "processed_data = process_data(data)", + ] + ) + header, sections, section_names = split_code_sections(code) + self.assertEqual( + header, + S( + [ + "# This is the main function", + "setup_workspace()", + ] + ), + ) + self.assertListEqual( + sections, + [ + S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + ] + ), + S( + [ + 'print("Section: Data Processing")', + "# Process data", + "processed_data = process_data(data)", + ] + ), + ], + ) + self.assertListEqual(section_names, ["Data Loading", "Data Processing"]) + + def test_happy_path_no_header(self): + code = S( + [ + "print('Section: Setup')", + "# This is the main function", + "setup_workspace()", + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ) + header, sections, section_names = split_code_sections(code) + self.assertEqual(header, None) + self.assertListEqual( + sections, + [ + S( + [ + "print('Section: Setup')", + "# This is the main function", + "setup_workspace()", + ] + ), + S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + ] + ), + S( + [ + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ), + ], + ) + self.assertListEqual(section_names, ["Setup", "Data Loading", "Data Processing"]) + + def test_wrong_format(self): + code = S( + [ + "# This is the main function", + "setup_workspace()", + "print('A Section: Data Loading')", + "# Load data", + "data = load_data()", + 's = """print(\'Section: Data Processing\')"""', + "# Process data", + "processed_data = process_data(data)", + ] + ) + header, sections, section_names = split_code_sections(code) + self.assertEqual( + header, + S( + [ + "# This is the main function", + "setup_workspace()", + "print('A Section: Data Loading')", + "# Load data", + "data = load_data()", + 's = """print(\'Section: Data Processing\')"""', + "# Process data", + "processed_data = process_data(data)", + ] + ), + ) + self.assertListEqual(sections, []) + self.assertListEqual(section_names, []) + + def test_empty(self): + code = "" + header, sections, section_names = split_code_sections(code) + self.assertEqual(header, None) + self.assertListEqual(sections, []) + self.assertListEqual(section_names, []) + + def test_single_no_sections(self): + code = "print('foo')" + header, sections, section_names = split_code_sections(code) + self.assertEqual(header, "print('foo')") + self.assertListEqual(sections, []) + self.assertListEqual(section_names, []) + + def test_single_with_section(self): + code = "print('Section: foo')" + header, sections, section_names = split_code_sections(code) + self.assertEqual(header, None) + self.assertListEqual(sections, ["print('Section: foo')"]) + self.assertListEqual(section_names, ["foo"]) + + def test_no_sections(self): + code = S( + [ + "# This is the main function", + "setup_workspace()", + "# Load data", + "data = load_data()", + "# Process data", + "processed_data = process_data(data)", + ] + ) + header, sections, section_names = split_code_sections(code) + self.assertEqual( + header, + S( + [ + "# This is the main function", + "setup_workspace()", + "# Load data", + "data = load_data()", + "# Process data", + "processed_data = process_data(data)", + ] + ), + ) + self.assertListEqual(sections, []) + self.assertListEqual(section_names, []) + + def test_ignores_indented_calls(self): + code = S( + [ + "# This is the main function", + "setup_workspace()", + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + "if some_condition():", + ' print("Section: Data Processing")', + " # Process data", + " processed_data = process_data(data)", + "", + "def print_section():", + " print('Section: Another Section')", + "", + "print('Section: Finalization')", + "# Finalize", + "finalize()", + ] + ) + header, sections, section_names = split_code_sections(code) + self.assertEqual( + header, + S( + [ + "# This is the main function", + "setup_workspace()", + ] + ), + ) + self.assertListEqual( + sections, + [ + S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + "if some_condition():", + ' print("Section: Data Processing")', + " # Process data", + " processed_data = process_data(data)", + "", + "def print_section():", + " print('Section: Another Section')", + "", + ] + ), + S(["print('Section: Finalization')", "# Finalize", "finalize()"]), + ], + ) + self.assertListEqual(section_names, ["Data Loading", "Finalization"]) + + +class TestSplitOutputSections(unittest.TestCase): + def test_happy_path(self): + output = S( + [ + "Setting up workspace...", + "Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ) + header, sections = split_output_sections(output, known_sections=["Data Loading", "Data Processing"]) + self.assertEqual( + header, + S( + [ + "Setting up workspace...", + ] + ), + ) + self.assertListEqual( + sections, + [ + S(["Section: Data Loading", "Loading data..."]), + S( + [ + "Section: Data Processing", + "Processing data...", + ] + ), + ], + ) + + def test_happy_path_no_header(self): + output = S( + [ + "Section: Setup", + "Setting up workspace...", + "Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ) + header, sections = split_output_sections(output, known_sections=["Setup", "Data Loading", "Data Processing"]) + self.assertEqual(header, None) + self.assertListEqual( + sections, + [ + S( + [ + "Section: Setup", + "Setting up workspace...", + ] + ), + S(["Section: Data Loading", "Loading data..."]), + S( + [ + "Section: Data Processing", + "Processing data...", + ] + ), + ], + ) + + def test_wrong_format(self): + output = S( + [ + "Setting up workspace...", + "Wrong Section: Data Loading", + "Loading data...", + "Wrong Section: Data Processing", + "Processing data...", + ] + ) + header, sections = split_output_sections(output, known_sections=["Data Loading", "Data Processing"]) + self.assertEqual( + header, + S( + [ + "Setting up workspace...", + "Wrong Section: Data Loading", + "Loading data...", + "Wrong Section: Data Processing", + "Processing data...", + ] + ), + ) + self.assertListEqual(sections, []) + + def test_empty(self): + output = "" + header, sections = split_output_sections(output, known_sections=["Data Loading", "Data Processing"]) + self.assertEqual(header, None) + self.assertListEqual(sections, []) + + def test_single_no_sections(self): + output = "foo" + header, sections = split_output_sections(output, known_sections=["foo"]) + self.assertEqual(header, "foo") + self.assertListEqual(sections, []) + + def test_single_with_section(self): + output = "Section: foo" + header, sections = split_output_sections(output, known_sections=["foo"]) + self.assertEqual(header, None) + self.assertListEqual(sections, ["Section: foo"]) + + def test_no_sections(self): + output = S( + [ + "Setting up workspace...", + "Loading data...", + "Processing data...", + ] + ) + header, sections = split_output_sections(output, known_sections=["Data Loading", "Data Processing"]) + self.assertEqual( + header, + S( + [ + "Setting up workspace...", + "Loading data...", + "Processing data...", + ] + ), + ) + self.assertListEqual(sections, []) + + def test_ignore_spaces(self): + output = S( + [ + "Setting up workspace...", + " Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ) + header, sections = split_output_sections(output, known_sections=["Data Loading", "Data Processing"]) + self.assertEqual( + header, + S( + [ + "Setting up workspace...", + " Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ), + ) + self.assertListEqual(sections, []) + + def test_ignore_unknown_section(self): + output = S( + [ + "Setting up workspace...", + "Section: Data Loading (1/5)", + "Section: Data Loading (2/5)", + "Section: Data Loading (3/5)", + "Section: Data Loading (4/5)", + "Section: Data Loading (5/5)", + "Loading data...", + "Section: Data Processing", + "Section: Data Processing (Sub task)", + "Processing data...", + ] + ) + header, sections = split_output_sections(output, known_sections=["Data Processing"]) + self.assertEqual( + header, + S( + [ + "Setting up workspace...", + "Section: Data Loading (1/5)", + "Section: Data Loading (2/5)", + "Section: Data Loading (3/5)", + "Section: Data Loading (4/5)", + "Section: Data Loading (5/5)", + "Loading data...", + ] + ), + ) + self.assertListEqual( + sections, + [ + S( + [ + "Section: Data Processing", + "Section: Data Processing (Sub task)", + "Processing data...", + ] + ), + ], + ) + + +class TestExtractSectionComments(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ) + comments, cleaned = extract_comment_under_first_print(code) + self.assertEqual(comments, "Load data") + self.assertEqual( + cleaned, + S( + [ + "print('Section: Data Loading')", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ), + ) + + def test_happy_path_multiline(self): + code = S( + [ + "print('Section: Data Loading')", + "# Load data", + "# This section loads some data", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ) + comments, cleaned = extract_comment_under_first_print(code) + self.assertEqual(comments, S(["Load data", "This section loads some data"])) + self.assertEqual( + cleaned, + S( + [ + "print('Section: Data Loading')", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ), + ) + + def test_no_comment(self): + code = S( + [ + "print('Section: Data Loading')", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ) + comments, cleaned = extract_comment_under_first_print(code) + self.assertEqual(comments, None) + self.assertEqual( + cleaned, + S( + [ + "print('Section: Data Loading')", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ), + ) + + def test_arbitrary_print_happy_path(self): + code = S( + [ + "print('No section here')", + "# Just a comment", + "data = load_data()", + ] + ) + comments, cleaned = extract_comment_under_first_print(code) + self.assertEqual(comments, "Just a comment") + self.assertEqual( + cleaned, + S( + [ + "print('No section here')", + "data = load_data()", + ] + ), + ) + + def test_empty_string(self): + code = "" + comments, cleaned = extract_comment_under_first_print(code) + self.assertEqual(comments, None) + self.assertEqual(cleaned, "") + + +class TestExtractFirstSectionNameFromCode(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "print('Section: Data Loading')", + "# Load data", + "data = load_data()", + "print('Section: Data Processing')", + "# Process data", + "processed_data = process_data(data)", + ] + ) + section_name = extract_first_section_name_from_code(code) + self.assertEqual(section_name, "Data Loading") + + def test_no_section(self): + code = S( + [ + "print('No section here')", + "# Just a comment", + "data = load_data()", + ] + ) + section_name = extract_first_section_name_from_code(code) + self.assertEqual(section_name, None) + + def test_empty_string(self): + code = "" + section_name = extract_first_section_name_from_code(code) + self.assertEqual(section_name, None) + + +class TestExtractFirstSectionNameFromOutput(unittest.TestCase): + def test_happy_path(self): + output = S( + [ + "Setting up workspace...", + "Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ) + section_name = extract_first_section_name_from_output(output) + self.assertEqual(section_name, "Data Loading") + + def test_no_section(self): + output = S( + [ + "Setting up workspace...", + "Loading data...", + "Processing data...", + ] + ) + section_name = extract_first_section_name_from_output(output) + self.assertEqual(section_name, None) + + def test_empty_string(self): + output = "" + section_name = extract_first_section_name_from_output(output) + self.assertEqual(section_name, None) + + +class TestIsFunctionCalled(unittest.TestCase): + def test_happy_path(self): + code = S(["def main():", " print('Hello World')", "", "main()"]) + self.assertTrue(is_function_called(code, "main")) + + def test_happy_path_with_args(self): + code = S( + [ + "main(123, 'abc')", + ] + ) + self.assertTrue(is_function_called(code, "main")) + + def test_happy_path_with_args_multiline(self): + code = S( + [ + "main(", + " 123,", + " 'abc'", + ")", + ] + ) + self.assertTrue(is_function_called(code, "main")) + + def test_not_called(self): + code = S( + [ + "def main():", + " print('Hello World')", + "", + ] + ) + self.assertFalse(is_function_called(code, "main")) + + def test_wrong_format(self): + code = S(["def main():", " print('Hello World')", "", "main2()"]) + self.assertFalse(is_function_called(code, "main")) + + def test_empty_string(self): + code = "" + self.assertFalse(is_function_called(code, "main")) + + +class TestRemoveFunction(unittest.TestCase): + def test_happy_path(self): + code = S(["def main():", " print('Hello World')", "", "main()"]) + cleaned_code = remove_function(code, "main") + expected_code = S(["", "main()"]) + self.assertEqual(cleaned_code, expected_code) + + def test_function_does_not_exist(self): + code = S(["def main2():", " print('Hello World')", "", "main()"]) + cleaned_code = remove_function(code, "main") + expected_code = S(["def main2():", " print('Hello World')", "", "main()"]) + self.assertEqual(cleaned_code, expected_code) + + def test_empty(self): + code = "" + cleaned_code = remove_function(code, "main") + expected_code = "" + self.assertEqual(cleaned_code, expected_code) + + def test_preserves_comments(self): + code = S( + [ + "def main():", + ' """' " This is the main function.", + ' """', + " print('Hello World')", + "", + "def main2():", + ' """' " This is the second main function.", + ' """', + " print('Hello World')", + "", + "# Some comment", + "main()", + ] + ) + cleaned_code = remove_function(code, "main") + expected_code = S( + [ + "", + "def main2():", + ' """' " This is the second main function.", + ' """', + " print('Hello World')", + "", + "# Some comment", + "main()", + ] + ) + self.assertEqual(cleaned_code, expected_code) + + +class TestRemoveMainBlock(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "if __name__ == '__main__':", + " main()", + ] + ) + cleaned_code = remove_main_block(code) + expected_code = "" + self.assertEqual(cleaned_code, expected_code) + + def test_one_liner(self): + code = S( + [ + "if __name__ == '__main__': main()", + ] + ) + cleaned_code = remove_main_block(code) + expected_code = "" + self.assertEqual(cleaned_code, expected_code) + + def test_happy_path_arbitrary_content(self): + code = S( + [ + "if __name__ == '__main__':", + " # foo", + " print('Hello World')", + " main()", + ] + ) + cleaned_code = remove_main_block(code) + expected_code = "" + self.assertEqual(cleaned_code, expected_code) + + def test_block_does_not_exist(self): + code = S( + [ + "if __name__ == '__foo__':", + " main()", + ] + ) + cleaned_code = remove_main_block(code) + expected_code = S( + [ + "if __name__ == '__foo__':", + " main()", + ] + ) + self.assertEqual(cleaned_code, expected_code) + + def test_empty(self): + code = "" + cleaned_code = remove_main_block(code) + expected_code = "" + self.assertEqual(cleaned_code, expected_code) + + +class TestExtractTopLevelFunctions(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "# This is the main function", + "", + "# Some more comments", + "def foo():", + " print('Hello World')", + "", + "def bar():", + " print('Helper function')", + ] + ) + functions = extract_top_level_functions_with_decorators_and_comments(code) + expected_fns = [ + ( + "foo", + S( + [ + "# This is the main function", + "", + "# Some more comments", + "def foo():", + " print('Hello World')", + "", + ] + ), + ), + ( + "bar", + S( + [ + "", + "def bar():", + " print('Helper function')", + ] + ), + ), + ] + self.assertEqual(len(functions), 2) + for idx, (name, segment) in enumerate(functions): + expected_name, expected_segment = expected_fns[idx] + self.assertIn(name, expected_name, "Function name should match") + self.assertIn(segment, expected_segment, "Function segment should match") + + def test_empty(self): + code = "" + functions = extract_top_level_functions_with_decorators_and_comments(code) + self.assertEqual(len(functions), 0) + + def test_stop_at_code(self): + code = S( + [ + "# This is the main function", + "foo = 123", + "# Some more comments", + "def foo():", + " print('Hello World')", + "", + "def bar():", + " print('Helper function')", + ] + ) + functions = extract_top_level_functions_with_decorators_and_comments(code) + expected_fns = [ + ( + "foo", + S( + [ + "# Some more comments", + "def foo():", + " print('Hello World')", + "", + ] + ), + ), + ( + "bar", + S( + [ + "", + "def bar():", + " print('Helper function')", + ] + ), + ), + ] + self.assertEqual(len(functions), 2) + for idx, (name, segment) in enumerate(functions): + expected_name, expected_segment = expected_fns[idx] + self.assertIn(name, expected_name, "Function name should match") + self.assertIn(segment, expected_segment, "Function segment should match") + + def test_trailing_comment(self): + code = S( + [ + "# This is the main function", + "", + "# Some more comments", + "def foo():", + " print('Hello World') # trailing comment", + "", + "def bar():", + " print('Helper function')", + ] + ) + functions = extract_top_level_functions_with_decorators_and_comments(code) + expected_fns = [ + ( + "foo", + S( + [ + "# This is the main function", + "", + "# Some more comments", + "def foo():", + " print('Hello World') # trailing comment", + "", + ] + ), + ), + ( + "bar", + S( + [ + "", + "def bar():", + " print('Helper function')", + ] + ), + ), + ] + self.assertEqual(len(functions), 2) + for idx, (name, segment) in enumerate(functions): + expected_name, expected_segment = expected_fns[idx] + self.assertIn(name, expected_name, "Function name should match") + self.assertIn(segment, expected_segment, "Function segment should match") + + +class TestSplitCodeAndOutputIntoSections(unittest.TestCase): + def test_happy_path(self): + code = S( + [ + "# Some notebook comments", + "import pandas as pd", + "", + "RANDOM_SEED = 42", + "" "def setup():", + " print('Setting up workspace...')", + "", + "def load_data():", + " return []", + "", + "def process_data(data):", + " return data", + "", + "def main():", + " setup()", + " print('Section: Data Loading')", + " # Load data", + " data = load_data()", + "", + " print('Section: Data Processing')", + " # Process data", + " processed_data = process_data(data)", + ] + ) + output = S( + [ + "Setting up workspace...", + "Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ) + sections = split_code_and_output_into_sections(code=code, stdout=output) + self.assertEqual(len(sections), 3) + self.assertDictEqual( + sections[0], + { + "name": None, + "comments": None, + "code": S( + [ + "# Some notebook comments", + "import pandas as pd", + "", + "RANDOM_SEED = 42", + "" "def setup():", + " print('Setting up workspace...')", + "", + "setup()", + ] + ), + "output": S(["Setting up workspace..."]), + }, + ) + self.assertDictEqual( + sections[1], + { + "name": "Data Loading", + "comments": "Load data", + "code": S( + [ + "def load_data():", + " return []", + "", + "print('Section: Data Loading')", + "data = load_data()", + ] + ), + "output": S( + [ + "Section: Data Loading", + "Loading data...", + ] + ), + }, + ) + self.assertDictEqual( + sections[2], + { + "name": "Data Processing", + "comments": "Process data", + "code": S( + [ + "def process_data(data):", + " return data", + "", + "print('Section: Data Processing')", + "processed_data = process_data(data)", + ] + ), + "output": S( + [ + "Section: Data Processing", + "Processing data...", + ] + ), + }, + ) + + def test_empty_code(self): + code = "" + output = S( + [ + "Setting up workspace...", + "Section: Data Loading", + "Loading data...", + "Section: Data Processing", + "Processing data...", + ] + ) + sections = split_code_and_output_into_sections(code=code, stdout=output) + self.assertEqual(len(sections), 3) + self.assertDictEqual( + sections[0], + { + "name": None, + "comments": None, + "code": "", + "output": S( + [ + "Setting up workspace...", + ] + ), + }, + ) + self.assertDictEqual( + sections[1], + { + "name": "Data Loading", + "comments": None, + "code": None, + "output": S( + [ + "Section: Data Loading", + "Loading data...", + ] + ), + }, + ) + self.assertDictEqual( + sections[2], + { + "name": "Data Processing", + "comments": None, + "code": None, + "output": S( + [ + "Section: Data Processing", + "Processing data...", + ] + ), + }, + ) + + def test_empty_outputs(self): + code = S( + [ + "# Some notebook comments", + "import pandas as pd", + "", + "RANDOM_SEED = 42", + "" "def setup():", + " print('Setting up workspace...')", + "", + "def load_data():", + " return []", + "", + "def process_data(data):", + " return data", + "", + "def main():", + " setup()", + " print('Section: Data Loading')", + " # Load data", + " data = load_data()", + "", + " print('Section: Data Processing')", + " # Process data", + " processed_data = process_data(data)", + ] + ) + output = "" + sections = split_code_and_output_into_sections(code=code, stdout=output) + self.assertEqual(len(sections), 3) + self.assertDictEqual( + sections[0], + { + "name": None, + "comments": None, + "code": S( + [ + "# Some notebook comments", + "import pandas as pd", + "", + "RANDOM_SEED = 42", + "" "def setup():", + " print('Setting up workspace...')", + "", + "setup()", + ] + ), + "output": None, + }, + ) + self.assertDictEqual( + sections[1], + { + "name": "Data Loading", + "comments": "Load data", + "code": S( + [ + "def load_data():", + " return []", + "", + "print('Section: Data Loading')", + "data = load_data()", + ] + ), + "output": None, + }, + ) + self.assertDictEqual( + sections[2], + { + "name": "Data Processing", + "comments": "Process data", + "code": S( + [ + "def process_data(data):", + " return data", + "", + "print('Section: Data Processing')", + "processed_data = process_data(data)", + ] + ), + "output": None, + }, + ) + + def test_ignored_sections(self): + code = S( + [ + "# Some notebook comments", + "import pandas as pd", + "", + "RANDOM_SEED = 42", + "" "def setup():", + " print('Setting up workspace...')", + "", + "def load_data():", + " return []", + "", + "def process_data(data):", + " return data", + "", + "def main():", + " setup()", + " print('Section: Data Loading')", + " if some_condition():", + " print('Section: Data Loading (sub task)')", + " # Load data", + " data = load_data()", + "", + " print('Section: Data Processing')", + " # Process data", + " for i in range(3):", + " print(f'Section: Data Processing {i}')", + " processed_data = process_data(data)", + ] + ) + output = S( + [ + "Setting up workspace...", + "Section: Data Loading", + "Section: Data Loading (sub task)", + "Loading data...", + "Section: Data Processing", + "Section: Data Processing 0", + "Section: Data Processing 1", + "Section: Data Processing 2", + "Processing data...", + ] + ) + sections = split_code_and_output_into_sections(code=code, stdout=output) + self.assertEqual(len(sections), 3) + self.assertDictEqual( + sections[0], + { + "name": None, + "comments": None, + "code": S( + [ + "# Some notebook comments", + "import pandas as pd", + "", + "RANDOM_SEED = 42", + "" "def setup():", + " print('Setting up workspace...')", + "", + "setup()", + ] + ), + "output": S(["Setting up workspace..."]), + }, + ) + self.assertDictEqual( + sections[1], + { + "name": "Data Loading", + "comments": None, + "code": S( + [ + "def load_data():", + " return []", + "", + "print('Section: Data Loading')", + "if some_condition():", + " print('Section: Data Loading (sub task)')", + "# Load data", + "data = load_data()", + ] + ), + "output": S( + [ + "Section: Data Loading", + "Section: Data Loading (sub task)", + "Loading data...", + ] + ), + }, + ) + self.assertDictEqual( + sections[2], + { + "name": "Data Processing", + "comments": "Process data", + "code": S( + [ + "def process_data(data):", + " return data", + "", + "print('Section: Data Processing')", + "for i in range(3):", + " print(f'Section: Data Processing {i}')", + "processed_data = process_data(data)", + ] + ), + "output": S( + [ + "Section: Data Processing", + "Section: Data Processing 0", + "Section: Data Processing 1", + "Section: Data Processing 2", + "Processing data...", + ] + ), + }, + ) + + def test_complex(self): + self.maxDiff = None + with open(os.path.join(test_files_dir, "main.py"), "r") as f: + code = f.read() + output = "" + sections = split_code_and_output_into_sections(code=code, stdout=output) + sections = split_code_and_output_into_sections(code=code, stdout=output) + self.assertEqual(len(sections), 6) + + expected_sections = [ + { + "name": None, + "comments": None, + "output": None, + "code": """import os +import sys +import time +import random +import numpy as np +import pandas as pd + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader + +import timm +import albumentations as A +from albumentations.pytorch import ToTensorV2 + +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import roc_auc_score, confusion_matrix + +import cv2 +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument('--debug', action='store_true', help='Run in debug mode') +args = parser.parse_args() +DEBUG = args.debug + +SEED = 2024 +np.random.seed(SEED) +random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +TRAIN_DIR = './workspace_input/train/' +TEST_DIR = './workspace_input/test/' +TRAIN_CSV = './workspace_input/train.csv' +SAMPLE_SUB_PATH = './workspace_input/sample_submission.csv' +MODEL_DIR = 'models/' +os.makedirs(MODEL_DIR, exist_ok=True) + +class CactusDataset(Dataset): + def __init__(self, image_ids, labels=None, id2path=None, transforms=None): + self.image_ids = image_ids + self.labels = labels + self.id2path = id2path + self.transforms = transforms + + def __len__(self): + return len(self.image_ids) + + def __getitem__(self, idx): + img_id = self.image_ids[idx] + img_path = self.id2path[img_id] + image = cv2.imread(img_path) + if image is None: + raise RuntimeError(f"Cannot read image at {img_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + if self.transforms: + augmented = self.transforms(image=image) + image = augmented["image"] + if self.labels is not None: + label = self.labels[idx] + return image, label, img_id + else: + return image, img_id + +""", + }, + { + "name": "Data Loading and Preprocessing", + "comments": "This section loads the train and test data, performs EDA, and prepares the dataset.", + "output": None, + "code": """def compute_class_weight(y): + counts = np.bincount(y) + if len(counts) < 2: + counts = np.pad(counts, (0, 2-len(counts)), constant_values=0) + n_pos, n_neg = counts[1], counts[0] + total = n_pos + n_neg + minority, majority = min(n_pos, n_neg), max(n_pos, n_neg) + ratio = majority / (minority + 1e-10) + need_weights = ratio > 2 + weights = None + if need_weights: + inv_freq = [1 / (n_neg + 1e-10), 1 / (n_pos + 1e-10)] + s = sum(inv_freq) + weights = [w / s * 2 for w in inv_freq] + return weights, n_pos, n_neg, ratio, need_weights + +def print_eda(train_df): + print("=== Start of EDA part ===") + print("Shape of train.csv:", train_df.shape) + print("First 5 rows:\\n", train_df.head()) + print("Column data types:\\n", train_df.dtypes) + print("Missing values per column:\\n", train_df.isnull().sum()) + print("Unique values per column:") + for col in train_df.columns: + print(f" - {col}: {train_df[col].nunique()}") + label_counts = train_df['has_cactus'].value_counts() + print("Label distribution (has_cactus):") + print(label_counts) + pos, neg = label_counts.get(1, 0), label_counts.get(0, 0) + total = pos + neg + if total > 0: + print(f" Positive:Negative ratio: {pos}:{neg} ({pos/total:.3f}:{neg/total:.3f})") + print(f" Percentage positive: {pos/total*100:.2f}%") + else: + print(" No data found.") + print("Image filename examples:", train_df['id'].unique()[:5]) + print("=== End of EDA part ===") + +print("Section: Data Loading and Preprocessing") +try: + train_df = pd.read_csv(TRAIN_CSV) +except Exception as e: + print(f"Failed to load train.csv: {e}") + sys.exit(1) +print_eda(train_df) + +train_id2path = {img_id: os.path.join(TRAIN_DIR, img_id) for img_id in train_df['id']} +try: + sample_sub = pd.read_csv(SAMPLE_SUB_PATH) +except Exception as e: + print(f"Failed to load sample_submission.csv: {e}") + sys.exit(1) +test_img_ids = list(sample_sub['id']) +test_id2path = {img_id: os.path.join(TEST_DIR, img_id) for img_id in test_img_ids} +print(f"Loaded {len(train_id2path)} train images, {len(test_id2path)} test images.") + +y_train = train_df['has_cactus'].values +class_weights, n_pos, n_neg, imbalance_ratio, need_weights = compute_class_weight(y_train) +print(f"Class stats: Pos={n_pos}, Neg={n_neg}, Imbalance Ratio(majority/minority)={imbalance_ratio:.3f}") +print(f"Use class weights: {need_weights}, Class weights: {class_weights if class_weights is not None else '[1.0,1.0]'}") +if class_weights is not None: + np.save(os.path.join(MODEL_DIR, "class_weights.npy"), class_weights)""", + }, + { + "name": "Feature Engineering", + "comments": None, + "output": None, + "code": """print("Section: Feature Engineering") +train_df = train_df.copy() +cv_fold = 5 +skf = StratifiedKFold(n_splits=cv_fold, shuffle=True, random_state=SEED) +folds = np.zeros(len(train_df), dtype=np.int32) +for idx, (_, val_idx) in enumerate(skf.split(train_df['id'], train_df['has_cactus'])): + folds[val_idx] = idx +train_df['fold'] = folds +print(f"Assigned stratified {cv_fold}-fold indices. Fold sample counts:") +for f in range(cv_fold): + dist = train_df.loc[train_df['fold'] == f, 'has_cactus'].value_counts().to_dict() + print(f" Fold {f}: n={len(train_df[train_df['fold'] == f])} class dist={dist}")""", + }, + { + "name": "Model Training and Evaluation", + "comments": None, + "output": None, + "code": """def inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, class_weights, need_weights, + BATCH_SIZE, N_WORKERS, cv_fold): + oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] + for fold in range(cv_fold): + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + val_ds = CactusDataset(val_img_ids, val_labels, id2path=train_id2path, transforms=get_transforms("val")) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(val_auc) + print(f"Reloaded fold {fold}, OOF Validation AUC={val_auc:.5f}") + + all_oof_true = np.concatenate(oof_true) + all_oof_pred = np.concatenate(oof_pred) + oof_auc = roc_auc_score(all_oof_true, all_oof_pred) + oof_cm = confusion_info(all_oof_true, all_oof_pred) + print(f"OOF ROC-AUC (from loaded models): {oof_auc:.5f}") + print(f"OOF Confusion Matrix:\\n{oof_cm}") + + test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") + ) + test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + test_pred_list = [] + for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") + test_probs = np.mean(test_pred_list, axis=0) + + submission = pd.read_csv(SAMPLE_SUB_PATH) + submission['has_cactus'] = test_probs + submission.to_csv('submission.csv', index=False) + print(f"Saved submission.csv in required format with {len(submission)} rows.") + + scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] + }) + scores_df.set_index('Model', inplace=True) + scores_df.to_csv("scores.csv") + print(f"Saved cross-validation scores to scores.csv") + +def confusion_info(y_true, y_pred, threshold=0.5): + preds = (y_pred > threshold).astype(int) + cm = confusion_matrix(y_true, preds) + return cm + +@torch.no_grad() +def eval_model(model, loss_fn, dataloader, device, class_weights): + model.eval() + y_true, y_pred = [], [] + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + probs = torch.sigmoid(logits) + y_true.append(labels.cpu().numpy()) + y_pred.append(probs.cpu().numpy()) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + y_true = np.vstack(y_true).reshape(-1) + y_pred = np.vstack(y_pred).reshape(-1) + avg_loss = total_loss / total_samples + return avg_loss, y_true, y_pred + +def train_one_epoch(model, loss_fn, optimizer, scheduler, dataloader, device, class_weights): + model.train() + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + if scheduler is not None: + scheduler.step() + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + avg_loss = total_loss / total_samples + return avg_loss + +def get_efficientnet_b3(dropout_rate=0.3): + model = timm.create_model('efficientnet_b3', pretrained=True) + n_in = model.classifier.in_features if hasattr(model, "classifier") else model.fc.in_features + model.classifier = nn.Sequential( + nn.Dropout(dropout_rate), + nn.Linear(n_in, 1) + ) + return model + +def get_dataloader(dataset, batch_size, shuffle=False, num_workers=4, pin_memory=True): + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=pin_memory + ) + +def get_transforms(mode='train'): + # Correct Cutout: Albumentations v1.4.15 provides 'Cutout' as a class, but not always in the root. + # Defensive import; fallback to the most robust method for v1.4.15 + imagenet_mean = [0.485, 0.456, 0.406] + imagenet_std = [0.229, 0.224, 0.225] + if mode == 'train': + min_frac, max_frac = 0.05, 0.2 + min_cut = int(300 * min_frac) + max_cut = int(300 * max_frac) + # There is no A.Cutout in v1.4.15 root, but A.augmentations.transforms.Cutout exists. + try: + from albumentations.augmentations.transforms import Cutout + have_cutout = True + except ImportError: + have_cutout = False + this_cut_h = random.randint(min_cut, max_cut) + this_cut_w = random.randint(min_cut, max_cut) + cutout_fill = [int(255 * m) for m in imagenet_mean] + tforms = [ + A.RandomResizedCrop(300, 300, scale=(0.7, 1.0), ratio=(0.8, 1.2), p=1.0), + A.Rotate(limit=30, p=0.8), + ] + if have_cutout: + tforms.append( + Cutout( + num_holes=1, + max_h_size=this_cut_h, + max_w_size=this_cut_w, + fill_value=cutout_fill, # RGB image in albumentations requires [R,G,B] + always_apply=False, + p=0.7 + ) + ) + else: + # No available Cutout, so fallback to no cutout but emit warning + print("WARNING: albumentations.Cutout not found, continuing without Cutout augmentation") + tforms.extend([ + A.RandomContrast(limit=0.2, p=0.5), + A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.1), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + return A.Compose(tforms) + else: + return A.Compose([ + A.Resize(300, 300), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + +print("Section: Model Training and Evaluation") +dropout_rate = round(random.uniform(0.2, 0.5), 2) +print(f"Model config: EfficientNet-B3, Image size 300, Head dropout={dropout_rate}") + +if DEBUG: + print("DEBUG mode: using 10% subsample and 1 epoch (per fold)") + sample_frac = 0.10 + sampled_idxs = [] + for f in range(cv_fold): + fold_idx = train_df.index[train_df['fold'] == f].tolist() + fold_labels = train_df.loc[fold_idx, 'has_cactus'].values + idx_pos = [i for i, l in zip(fold_idx, fold_labels) if l == 1] + idx_neg = [i for i, l in zip(fold_idx, fold_labels) if l == 0] + n_pos = max(1, int(sample_frac * len(idx_pos))) + n_neg = max(1, int(sample_frac * len(idx_neg))) + if len(idx_pos) > 0: + sampled_idxs += np.random.choice(idx_pos, n_pos, replace=False).tolist() + if len(idx_neg) > 0: + sampled_idxs += np.random.choice(idx_neg, n_neg, replace=False).tolist() + train_df = train_df.loc[sampled_idxs].reset_index(drop=True) + print(f"DEBUG subsample shape: {train_df.shape}") + debug_epochs = 1 +else: + debug_epochs = None + +BATCH_SIZE = 64 if torch.cuda.is_available() else 32 +N_WORKERS = 4 if torch.cuda.is_available() else 1 +EPOCHS = 20 if not DEBUG else debug_epochs +MIN_EPOCHS = 5 if not DEBUG else 1 +EARLY_STOP_PATIENCE = 7 if not DEBUG else 2 +LR = 1e-3 + +model_files = [os.path.join(MODEL_DIR, f"efficientnet_b3_fold{f}.pt") for f in range(cv_fold)] +if all([os.path.exists(f) for f in model_files]): + print("All fold models found in models/. Running inference and file saving only (no retrain).") + inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, + class_weights, need_weights, BATCH_SIZE, N_WORKERS, cv_fold) + return + +oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] +start_time = time.time() if DEBUG else None + +for fold in range(cv_fold): + print(f"\\n=== FOLD {fold} TRAINING ===") + df_train = train_df[train_df['fold'] != fold].reset_index(drop=True) + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + print(f"Train size: {df_train.shape[0]}, Val size: {df_val.shape[0]}") + train_img_ids = df_train['id'].tolist() + train_labels = df_train['has_cactus'].values + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + + train_ds = CactusDataset( + train_img_ids, train_labels, + id2path=train_id2path, + transforms=get_transforms("train") + ) + val_ds = CactusDataset( + val_img_ids, val_labels, + id2path=train_id2path, + transforms=get_transforms("val") + ) + train_loader = get_dataloader(train_ds, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + optimizer = optim.AdamW(model.parameters(), lr=LR) + scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + best_auc = -np.inf + best_epoch = -1 + best_model_state = None + patience = 0 + + for epoch in range(EPOCHS): + train_loss = train_one_epoch( + model, loss_fn, optimizer, scheduler, train_loader, DEVICE, fold_class_weights) + val_loss, val_true, val_pred = eval_model( + model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + cm = confusion_info(val_true, val_pred) + print(f"Epoch {epoch+1:02d}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_auc={val_auc:.4f}") + print(f" Val confusion_matrix (rows:true [0,1]; cols:pred [0,1]):\\n{cm}") + if val_auc > best_auc: + best_auc = val_auc + best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} + best_epoch = epoch + patience = 0 + else: + patience += 1 + if DEBUG and epoch + 1 >= debug_epochs: + break + if (epoch + 1) >= MIN_EPOCHS and patience >= EARLY_STOP_PATIENCE: + print(f"Early stopping at epoch {epoch+1}, best_epoch={best_epoch+1}.") + break + + model.load_state_dict(best_model_state) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + torch.save(model.state_dict(), fold_model_path) + print(f"Saved best model for fold {fold} at {fold_model_path} (best_auc={best_auc:.5f}, best_epoch={best_epoch+1})") + + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(best_auc) + print(f"OOF stored for fold {fold}, Validation AUC={best_auc:.5f}") + +end_time = time.time() if DEBUG else None +if DEBUG: + debug_time = end_time - start_time + estimated_time = (1 / 0.1) * (EPOCHS / debug_epochs) * debug_time + print("=== Start of Debug Information ===") + print(f"debug_time: {debug_time:.1f}") + print(f"estimated_time: {estimated_time:.1f}") + print("=== End of Debug Information ===")""", + }, + { + "name": "Ensemble Strategy and Final Predictions", + "comments": None, + "output": None, + "code": """print("Section: Ensemble Strategy and Final Predictions") +all_oof_true = np.concatenate(oof_true) +all_oof_pred = np.concatenate(oof_pred) +oof_auc = roc_auc_score(all_oof_true, all_oof_pred) +oof_cm = confusion_info(all_oof_true, all_oof_pred) +print(f"OOF ROC-AUC: {oof_auc:.5f}") +print(f"OOF Confusion Matrix:\\n{oof_cm}") + +test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") +) +test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) +test_pred_list = [] +for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") +test_probs = np.mean(test_pred_list, axis=0)""", + }, + { + "name": "Submission File Generation", + "comments": None, + "output": None, + "code": """print("Section: Submission File Generation") +submission = pd.read_csv(SAMPLE_SUB_PATH) +submission['has_cactus'] = test_probs +submission.to_csv('submission.csv', index=False) +print(f"Saved submission.csv in required format with {len(submission)} rows.") + +scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] +}) +scores_df.set_index('Model', inplace=True) +scores_df.to_csv("scores.csv") +print(f"Saved cross-validation scores to scores.csv")""", + }, + ] + + for i, section in enumerate(sections): + self.assertEqual( + section["name"], + expected_sections[i]["name"], + f"Section {i} name mismatch", + ) + self.assertEqual( + section["comments"], + expected_sections[i]["comments"], + f"Section {i} comments mismatch", + ) + self.assertEqual( + section["output"], + expected_sections[i]["output"], + f"Section {i} output mismatch", + ) + self.assertEqual( + section["code"], + expected_sections[i]["code"], + f"Section {i} code mismatch", + ) + + +def S(s_arr): + return "\n".join(s_arr) + + +if __name__ == "__main__": + unittest.main() + # pytest test/notebook/test_util.py diff --git a/test/notebook/testfiles/main.ipynb b/test/notebook/testfiles/main.ipynb new file mode 100644 index 00000000..f39198d6 --- /dev/null +++ b/test/notebook/testfiles/main.ipynb @@ -0,0 +1,608 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "ebeca6b7", + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "# hack to allow argparse to work in notebook\n", + "sys.argv = [\"main.py\"]\n", + "\n", + "import os\n", + "import time\n", + "import random\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from torch.utils.data import Dataset, DataLoader\n", + "\n", + "import timm\n", + "import albumentations as A\n", + "from albumentations.pytorch import ToTensorV2\n", + "\n", + "from sklearn.model_selection import StratifiedKFold\n", + "from sklearn.metrics import roc_auc_score, confusion_matrix\n", + "\n", + "import cv2\n", + "import argparse\n", + "\n", + "parser = argparse.ArgumentParser()\n", + "parser.add_argument('--debug', action='store_true', help='Run in debug mode')\n", + "args = parser.parse_args()\n", + "DEBUG = args.debug\n", + "\n", + "SEED = 2024\n", + "np.random.seed(SEED)\n", + "random.seed(SEED)\n", + "torch.manual_seed(SEED)\n", + "torch.cuda.manual_seed_all(SEED)\n", + "\n", + "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", + "TRAIN_DIR = './workspace_input/train/'\n", + "TEST_DIR = './workspace_input/test/'\n", + "TRAIN_CSV = './workspace_input/train.csv'\n", + "SAMPLE_SUB_PATH = './workspace_input/sample_submission.csv'\n", + "MODEL_DIR = 'models/'\n", + "os.makedirs(MODEL_DIR, exist_ok=True)\n", + "\n", + "class CactusDataset(Dataset):\n", + " def __init__(self, image_ids, labels=None, id2path=None, transforms=None):\n", + " self.image_ids = image_ids\n", + " self.labels = labels\n", + " self.id2path = id2path\n", + " self.transforms = transforms\n", + "\n", + " def __len__(self):\n", + " return len(self.image_ids)\n", + "\n", + " def __getitem__(self, idx):\n", + " img_id = self.image_ids[idx]\n", + " img_path = self.id2path[img_id]\n", + " image = cv2.imread(img_path)\n", + " if image is None:\n", + " raise RuntimeError(f\"Cannot read image at {img_path}\")\n", + " image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n", + " if self.transforms:\n", + " augmented = self.transforms(image=image)\n", + " image = augmented[\"image\"]\n", + " if self.labels is not None:\n", + " label = self.labels[idx]\n", + " return image, label, img_id\n", + " else:\n", + " return image, img_id\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "9086e8dc", + "metadata": {}, + "source": [ + "## Data Loading and Preprocessing\n", + "This section loads the train and test data, performs EDA, and prepares the dataset.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "05509a31", + "metadata": {}, + "outputs": [], + "source": [ + "def compute_class_weight(y):\n", + " counts = np.bincount(y)\n", + " if len(counts) < 2:\n", + " counts = np.pad(counts, (0, 2-len(counts)), constant_values=0)\n", + " n_pos, n_neg = counts[1], counts[0]\n", + " total = n_pos + n_neg\n", + " minority, majority = min(n_pos, n_neg), max(n_pos, n_neg)\n", + " ratio = majority / (minority + 1e-10)\n", + " need_weights = ratio > 2\n", + " weights = None\n", + " if need_weights:\n", + " inv_freq = [1 / (n_neg + 1e-10), 1 / (n_pos + 1e-10)]\n", + " s = sum(inv_freq)\n", + " weights = [w / s * 2 for w in inv_freq]\n", + " return weights, n_pos, n_neg, ratio, need_weights\n", + "\n", + "def print_eda(train_df):\n", + " print(\"=== Start of EDA part ===\")\n", + " print(\"Shape of train.csv:\", train_df.shape)\n", + " print(\"First 5 rows:\\n\", train_df.head())\n", + " print(\"Column data types:\\n\", train_df.dtypes)\n", + " print(\"Missing values per column:\\n\", train_df.isnull().sum())\n", + " print(\"Unique values per column:\")\n", + " for col in train_df.columns:\n", + " print(f\" - {col}: {train_df[col].nunique()}\")\n", + " label_counts = train_df['has_cactus'].value_counts()\n", + " print(\"Label distribution (has_cactus):\")\n", + " print(label_counts)\n", + " pos, neg = label_counts.get(1, 0), label_counts.get(0, 0)\n", + " total = pos + neg\n", + " if total > 0:\n", + " print(f\" Positive:Negative ratio: {pos}:{neg} ({pos/total:.3f}:{neg/total:.3f})\")\n", + " print(f\" Percentage positive: {pos/total*100:.2f}%\")\n", + " else:\n", + " print(\" No data found.\")\n", + " print(\"Image filename examples:\", train_df['id'].unique()[:5])\n", + " print(\"=== End of EDA part ===\")\n", + "\n", + "print(\"Section: Data Loading and Preprocessing\")\n", + "try:\n", + " train_df = pd.read_csv(TRAIN_CSV)\n", + "except Exception as e:\n", + " print(f\"Failed to load train.csv: {e}\")\n", + " sys.exit(1)\n", + "print_eda(train_df)\n", + "\n", + "train_id2path = {img_id: os.path.join(TRAIN_DIR, img_id) for img_id in train_df['id']}\n", + "try:\n", + " sample_sub = pd.read_csv(SAMPLE_SUB_PATH)\n", + "except Exception as e:\n", + " print(f\"Failed to load sample_submission.csv: {e}\")\n", + " sys.exit(1)\n", + "test_img_ids = list(sample_sub['id'])\n", + "test_id2path = {img_id: os.path.join(TEST_DIR, img_id) for img_id in test_img_ids}\n", + "print(f\"Loaded {len(train_id2path)} train images, {len(test_id2path)} test images.\")\n", + "\n", + "y_train = train_df['has_cactus'].values\n", + "class_weights, n_pos, n_neg, imbalance_ratio, need_weights = compute_class_weight(y_train)\n", + "print(f\"Class stats: Pos={n_pos}, Neg={n_neg}, Imbalance Ratio(majority/minority)={imbalance_ratio:.3f}\")\n", + "print(f\"Use class weights: {need_weights}, Class weights: {class_weights if class_weights is not None else '[1.0,1.0]'}\")\n", + "if class_weights is not None:\n", + " np.save(os.path.join(MODEL_DIR, \"class_weights.npy\"), class_weights)" + ] + }, + { + "cell_type": "markdown", + "id": "b201cd3f", + "metadata": {}, + "source": [ + "## Feature Engineering\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d7d4697e", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Feature Engineering\")\n", + "train_df = train_df.copy()\n", + "cv_fold = 5\n", + "skf = StratifiedKFold(n_splits=cv_fold, shuffle=True, random_state=SEED)\n", + "folds = np.zeros(len(train_df), dtype=np.int32)\n", + "for idx, (_, val_idx) in enumerate(skf.split(train_df['id'], train_df['has_cactus'])):\n", + " folds[val_idx] = idx\n", + "train_df['fold'] = folds\n", + "print(f\"Assigned stratified {cv_fold}-fold indices. Fold sample counts:\")\n", + "for f in range(cv_fold):\n", + " dist = train_df.loc[train_df['fold'] == f, 'has_cactus'].value_counts().to_dict()\n", + " print(f\" Fold {f}: n={len(train_df[train_df['fold'] == f])} class dist={dist}\")" + ] + }, + { + "cell_type": "markdown", + "id": "23e606da", + "metadata": {}, + "source": [ + "## Model Training and Evaluation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "853b0c24", + "metadata": {}, + "outputs": [], + "source": [ + "def inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, class_weights, need_weights,\n", + " BATCH_SIZE, N_WORKERS, cv_fold):\n", + " oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], []\n", + " for fold in range(cv_fold):\n", + " df_val = train_df[train_df['fold'] == fold].reset_index(drop=True)\n", + " val_img_ids = df_val['id'].tolist()\n", + " val_labels = df_val['has_cactus'].values\n", + " val_ds = CactusDataset(val_img_ids, val_labels, id2path=train_id2path, transforms=get_transforms(\"val\"))\n", + " val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS)\n", + " fold_model_path = os.path.join(MODEL_DIR, f\"efficientnet_b3_fold{fold}.pt\")\n", + " model = get_efficientnet_b3(dropout_rate=dropout_rate)\n", + " model.load_state_dict(torch.load(fold_model_path, map_location='cpu'))\n", + " model.to(DEVICE)\n", + " model.eval()\n", + " fold_class_weights = class_weights if need_weights else None\n", + " if fold_class_weights is not None:\n", + " fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE)\n", + " loss_fn = nn.BCEWithLogitsLoss(reduction='none')\n", + " _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights)\n", + " val_auc = roc_auc_score(val_true, val_pred)\n", + " oof_true.append(val_true)\n", + " oof_pred.append(val_pred)\n", + " fold_val_ids.append(val_img_ids)\n", + " fold_scores.append(val_auc)\n", + " print(f\"Reloaded fold {fold}, OOF Validation AUC={val_auc:.5f}\")\n", + "\n", + " all_oof_true = np.concatenate(oof_true)\n", + " all_oof_pred = np.concatenate(oof_pred)\n", + " oof_auc = roc_auc_score(all_oof_true, all_oof_pred)\n", + " oof_cm = confusion_info(all_oof_true, all_oof_pred)\n", + " print(f\"OOF ROC-AUC (from loaded models): {oof_auc:.5f}\")\n", + " print(f\"OOF Confusion Matrix:\\n{oof_cm}\")\n", + "\n", + " test_ds = CactusDataset(\n", + " test_img_ids, labels=None,\n", + " id2path=test_id2path,\n", + " transforms=get_transforms(\"val\")\n", + " )\n", + " test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS)\n", + " test_pred_list = []\n", + " for fold in range(cv_fold):\n", + " fold_model_path = os.path.join(MODEL_DIR, f\"efficientnet_b3_fold{fold}.pt\")\n", + " model = get_efficientnet_b3(dropout_rate=dropout_rate)\n", + " model.load_state_dict(torch.load(fold_model_path, map_location='cpu'))\n", + " model.to(DEVICE)\n", + " model.eval()\n", + " preds = []\n", + " with torch.no_grad():\n", + " for batch in test_loader:\n", + " images, img_ids = batch\n", + " images = images.to(DEVICE)\n", + " logits = model(images)\n", + " probs = torch.sigmoid(logits).cpu().numpy().reshape(-1)\n", + " preds.append(probs)\n", + " fold_test_pred = np.concatenate(preds)\n", + " test_pred_list.append(fold_test_pred)\n", + " print(f\"Loaded fold {fold} for test prediction.\")\n", + " test_probs = np.mean(test_pred_list, axis=0)\n", + "\n", + " submission = pd.read_csv(SAMPLE_SUB_PATH)\n", + " submission['has_cactus'] = test_probs\n", + " submission.to_csv('submission.csv', index=False)\n", + " print(f\"Saved submission.csv in required format with {len(submission)} rows.\")\n", + "\n", + " scores_df = pd.DataFrame({\n", + " 'Model': [f\"efficientnet_b3_fold{f}\" for f in range(cv_fold)] + ['ensemble'],\n", + " 'ROC-AUC': list(fold_scores) + [oof_auc]\n", + " })\n", + " scores_df.set_index('Model', inplace=True)\n", + " scores_df.to_csv(\"scores.csv\")\n", + " print(f\"Saved cross-validation scores to scores.csv\")\n", + "\n", + "def confusion_info(y_true, y_pred, threshold=0.5):\n", + " preds = (y_pred > threshold).astype(int)\n", + " cm = confusion_matrix(y_true, preds)\n", + " return cm\n", + "\n", + "@torch.no_grad()\n", + "def eval_model(model, loss_fn, dataloader, device, class_weights):\n", + " model.eval()\n", + " y_true, y_pred = [], []\n", + " total_loss = 0.0\n", + " total_samples = 0\n", + " for batch in dataloader:\n", + " images, labels, _ = batch\n", + " images = images.to(device)\n", + " labels = labels.float().unsqueeze(1).to(device)\n", + " logits = model(images)\n", + " probs = torch.sigmoid(logits)\n", + " y_true.append(labels.cpu().numpy())\n", + " y_pred.append(probs.cpu().numpy())\n", + " if class_weights is not None:\n", + " weight = labels * class_weights[1] + (1 - labels) * class_weights[0]\n", + " loss = loss_fn(logits, labels)\n", + " loss = (loss * weight).mean()\n", + " else:\n", + " loss = loss_fn(logits, labels)\n", + " total_loss += loss.item() * labels.size(0)\n", + " total_samples += labels.size(0)\n", + " y_true = np.vstack(y_true).reshape(-1)\n", + " y_pred = np.vstack(y_pred).reshape(-1)\n", + " avg_loss = total_loss / total_samples\n", + " return avg_loss, y_true, y_pred\n", + "\n", + "def train_one_epoch(model, loss_fn, optimizer, scheduler, dataloader, device, class_weights):\n", + " model.train()\n", + " total_loss = 0.0\n", + " total_samples = 0\n", + " for batch in dataloader:\n", + " images, labels, _ = batch\n", + " images = images.to(device)\n", + " labels = labels.float().unsqueeze(1).to(device)\n", + " logits = model(images)\n", + " if class_weights is not None:\n", + " weight = labels * class_weights[1] + (1 - labels) * class_weights[0]\n", + " loss = loss_fn(logits, labels)\n", + " loss = (loss * weight).mean()\n", + " else:\n", + " loss = loss_fn(logits, labels)\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + " if scheduler is not None:\n", + " scheduler.step()\n", + " total_loss += loss.item() * labels.size(0)\n", + " total_samples += labels.size(0)\n", + " avg_loss = total_loss / total_samples\n", + " return avg_loss\n", + "\n", + "def get_efficientnet_b3(dropout_rate=0.3):\n", + " model = timm.create_model('efficientnet_b3', pretrained=True)\n", + " n_in = model.classifier.in_features if hasattr(model, \"classifier\") else model.fc.in_features\n", + " model.classifier = nn.Sequential(\n", + " nn.Dropout(dropout_rate),\n", + " nn.Linear(n_in, 1)\n", + " )\n", + " return model\n", + "\n", + "def get_dataloader(dataset, batch_size, shuffle=False, num_workers=4, pin_memory=True):\n", + " return DataLoader(\n", + " dataset,\n", + " batch_size=batch_size,\n", + " shuffle=shuffle,\n", + " num_workers=num_workers,\n", + " pin_memory=pin_memory\n", + " )\n", + "\n", + "def get_transforms(mode='train'):\n", + " # Correct Cutout: Albumentations v1.4.15 provides 'Cutout' as a class, but not always in the root.\n", + " # Defensive import; fallback to the most robust method for v1.4.15\n", + " imagenet_mean = [0.485, 0.456, 0.406]\n", + " imagenet_std = [0.229, 0.224, 0.225]\n", + " if mode == 'train':\n", + " min_frac, max_frac = 0.05, 0.2\n", + " min_cut = int(300 * min_frac)\n", + " max_cut = int(300 * max_frac)\n", + " # There is no A.Cutout in v1.4.15 root, but A.augmentations.transforms.Cutout exists.\n", + " try:\n", + " from albumentations.augmentations.transforms import Cutout\n", + " have_cutout = True\n", + " except ImportError:\n", + " have_cutout = False\n", + " this_cut_h = random.randint(min_cut, max_cut)\n", + " this_cut_w = random.randint(min_cut, max_cut)\n", + " cutout_fill = [int(255 * m) for m in imagenet_mean]\n", + " tforms = [\n", + " A.RandomResizedCrop(300, 300, scale=(0.7, 1.0), ratio=(0.8, 1.2), p=1.0),\n", + " A.Rotate(limit=30, p=0.8),\n", + " ]\n", + " if have_cutout:\n", + " tforms.append(\n", + " Cutout(\n", + " num_holes=1,\n", + " max_h_size=this_cut_h,\n", + " max_w_size=this_cut_w,\n", + " fill_value=cutout_fill, # RGB image in albumentations requires [R,G,B]\n", + " always_apply=False,\n", + " p=0.7\n", + " )\n", + " )\n", + " else:\n", + " # No available Cutout, so fallback to no cutout but emit warning\n", + " print(\"WARNING: albumentations.Cutout not found, continuing without Cutout augmentation\")\n", + " tforms.extend([\n", + " A.RandomContrast(limit=0.2, p=0.5),\n", + " A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.1),\n", + " A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0),\n", + " ToTensorV2()\n", + " ])\n", + " return A.Compose(tforms)\n", + " else:\n", + " return A.Compose([\n", + " A.Resize(300, 300),\n", + " A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0),\n", + " ToTensorV2()\n", + " ])\n", + "\n", + "print(\"Section: Model Training and Evaluation\")\n", + "dropout_rate = round(random.uniform(0.2, 0.5), 2)\n", + "print(f\"Model config: EfficientNet-B3, Image size 300, Head dropout={dropout_rate}\")\n", + "\n", + "if DEBUG:\n", + " print(\"DEBUG mode: using 10% subsample and 1 epoch (per fold)\")\n", + " sample_frac = 0.10\n", + " sampled_idxs = []\n", + " for f in range(cv_fold):\n", + " fold_idx = train_df.index[train_df['fold'] == f].tolist()\n", + " fold_labels = train_df.loc[fold_idx, 'has_cactus'].values\n", + " idx_pos = [i for i, l in zip(fold_idx, fold_labels) if l == 1]\n", + " idx_neg = [i for i, l in zip(fold_idx, fold_labels) if l == 0]\n", + " n_pos = max(1, int(sample_frac * len(idx_pos)))\n", + " n_neg = max(1, int(sample_frac * len(idx_neg)))\n", + " if len(idx_pos) > 0:\n", + " sampled_idxs += np.random.choice(idx_pos, n_pos, replace=False).tolist()\n", + " if len(idx_neg) > 0:\n", + " sampled_idxs += np.random.choice(idx_neg, n_neg, replace=False).tolist()\n", + " train_df = train_df.loc[sampled_idxs].reset_index(drop=True)\n", + " print(f\"DEBUG subsample shape: {train_df.shape}\")\n", + " debug_epochs = 1\n", + "else:\n", + " debug_epochs = None\n", + "\n", + "BATCH_SIZE = 64 if torch.cuda.is_available() else 32\n", + "N_WORKERS = 4 if torch.cuda.is_available() else 1\n", + "EPOCHS = 20 if not DEBUG else debug_epochs\n", + "MIN_EPOCHS = 5 if not DEBUG else 1\n", + "EARLY_STOP_PATIENCE = 7 if not DEBUG else 2\n", + "LR = 1e-3\n", + "\n", + "model_files = [os.path.join(MODEL_DIR, f\"efficientnet_b3_fold{f}.pt\") for f in range(cv_fold)]\n", + "if all([os.path.exists(f) for f in model_files]):\n", + " print(\"All fold models found in models/. Running inference and file saving only (no retrain).\")\n", + " inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate,\n", + " class_weights, need_weights, BATCH_SIZE, N_WORKERS, cv_fold)\n", + " return\n", + "\n", + "oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], []\n", + "start_time = time.time() if DEBUG else None\n", + "\n", + "for fold in range(cv_fold):\n", + " print(f\"\\n=== FOLD {fold} TRAINING ===\")\n", + " df_train = train_df[train_df['fold'] != fold].reset_index(drop=True)\n", + " df_val = train_df[train_df['fold'] == fold].reset_index(drop=True)\n", + " print(f\"Train size: {df_train.shape[0]}, Val size: {df_val.shape[0]}\")\n", + " train_img_ids = df_train['id'].tolist()\n", + " train_labels = df_train['has_cactus'].values\n", + " val_img_ids = df_val['id'].tolist()\n", + " val_labels = df_val['has_cactus'].values\n", + "\n", + " train_ds = CactusDataset(\n", + " train_img_ids, train_labels,\n", + " id2path=train_id2path,\n", + " transforms=get_transforms(\"train\")\n", + " )\n", + " val_ds = CactusDataset(\n", + " val_img_ids, val_labels,\n", + " id2path=train_id2path,\n", + " transforms=get_transforms(\"val\")\n", + " )\n", + " train_loader = get_dataloader(train_ds, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS)\n", + " val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS)\n", + " model = get_efficientnet_b3(dropout_rate=dropout_rate)\n", + " model.to(DEVICE)\n", + " loss_fn = nn.BCEWithLogitsLoss(reduction='none')\n", + " optimizer = optim.AdamW(model.parameters(), lr=LR)\n", + " scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)\n", + " fold_class_weights = class_weights if need_weights else None\n", + " if fold_class_weights is not None:\n", + " fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE)\n", + " best_auc = -np.inf\n", + " best_epoch = -1\n", + " best_model_state = None\n", + " patience = 0\n", + "\n", + " for epoch in range(EPOCHS):\n", + " train_loss = train_one_epoch(\n", + " model, loss_fn, optimizer, scheduler, train_loader, DEVICE, fold_class_weights)\n", + " val_loss, val_true, val_pred = eval_model(\n", + " model, loss_fn, val_loader, DEVICE, fold_class_weights)\n", + " val_auc = roc_auc_score(val_true, val_pred)\n", + " cm = confusion_info(val_true, val_pred)\n", + " print(f\"Epoch {epoch+1:02d}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_auc={val_auc:.4f}\")\n", + " print(f\" Val confusion_matrix (rows:true [0,1]; cols:pred [0,1]):\\n{cm}\")\n", + " if val_auc > best_auc:\n", + " best_auc = val_auc\n", + " best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}\n", + " best_epoch = epoch\n", + " patience = 0\n", + " else:\n", + " patience += 1\n", + " if DEBUG and epoch + 1 >= debug_epochs:\n", + " break\n", + " if (epoch + 1) >= MIN_EPOCHS and patience >= EARLY_STOP_PATIENCE:\n", + " print(f\"Early stopping at epoch {epoch+1}, best_epoch={best_epoch+1}.\")\n", + " break\n", + "\n", + " model.load_state_dict(best_model_state)\n", + " fold_model_path = os.path.join(MODEL_DIR, f\"efficientnet_b3_fold{fold}.pt\")\n", + " torch.save(model.state_dict(), fold_model_path)\n", + " print(f\"Saved best model for fold {fold} at {fold_model_path} (best_auc={best_auc:.5f}, best_epoch={best_epoch+1})\")\n", + "\n", + " _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights)\n", + " oof_true.append(val_true)\n", + " oof_pred.append(val_pred)\n", + " fold_val_ids.append(val_img_ids)\n", + " fold_scores.append(best_auc)\n", + " print(f\"OOF stored for fold {fold}, Validation AUC={best_auc:.5f}\")\n", + "\n", + "end_time = time.time() if DEBUG else None\n", + "if DEBUG:\n", + " debug_time = end_time - start_time\n", + " estimated_time = (1 / 0.1) * (EPOCHS / debug_epochs) * debug_time\n", + " print(\"=== Start of Debug Information ===\")\n", + " print(f\"debug_time: {debug_time:.1f}\")\n", + " print(f\"estimated_time: {estimated_time:.1f}\")\n", + " print(\"=== End of Debug Information ===\")" + ] + }, + { + "cell_type": "markdown", + "id": "c3f0269e", + "metadata": {}, + "source": [ + "## Ensemble Strategy and Final Predictions\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "308dcdb4", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Ensemble Strategy and Final Predictions\")\n", + "all_oof_true = np.concatenate(oof_true)\n", + "all_oof_pred = np.concatenate(oof_pred)\n", + "oof_auc = roc_auc_score(all_oof_true, all_oof_pred)\n", + "oof_cm = confusion_info(all_oof_true, all_oof_pred)\n", + "print(f\"OOF ROC-AUC: {oof_auc:.5f}\")\n", + "print(f\"OOF Confusion Matrix:\\n{oof_cm}\")\n", + "\n", + "test_ds = CactusDataset(\n", + " test_img_ids, labels=None,\n", + " id2path=test_id2path,\n", + " transforms=get_transforms(\"val\")\n", + ")\n", + "test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS)\n", + "test_pred_list = []\n", + "for fold in range(cv_fold):\n", + " fold_model_path = os.path.join(MODEL_DIR, f\"efficientnet_b3_fold{fold}.pt\")\n", + " model = get_efficientnet_b3(dropout_rate=dropout_rate)\n", + " model.load_state_dict(torch.load(fold_model_path, map_location='cpu'))\n", + " model.to(DEVICE)\n", + " model.eval()\n", + " preds = []\n", + " with torch.no_grad():\n", + " for batch in test_loader:\n", + " images, img_ids = batch\n", + " images = images.to(DEVICE)\n", + " logits = model(images)\n", + " probs = torch.sigmoid(logits).cpu().numpy().reshape(-1)\n", + " preds.append(probs)\n", + " fold_test_pred = np.concatenate(preds)\n", + " test_pred_list.append(fold_test_pred)\n", + " print(f\"Loaded fold {fold} for test prediction.\")\n", + "test_probs = np.mean(test_pred_list, axis=0)" + ] + }, + { + "cell_type": "markdown", + "id": "58b5ded8", + "metadata": {}, + "source": [ + "## Submission File Generation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "988914c8", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Submission File Generation\")\n", + "submission = pd.read_csv(SAMPLE_SUB_PATH)\n", + "submission['has_cactus'] = test_probs\n", + "submission.to_csv('submission.csv', index=False)\n", + "print(f\"Saved submission.csv in required format with {len(submission)} rows.\")\n", + "\n", + "scores_df = pd.DataFrame({\n", + " 'Model': [f\"efficientnet_b3_fold{f}\" for f in range(cv_fold)] + ['ensemble'],\n", + " 'ROC-AUC': list(fold_scores) + [oof_auc]\n", + "})\n", + "scores_df.set_index('Model', inplace=True)\n", + "scores_df.to_csv(\"scores.csv\")\n", + "print(f\"Saved cross-validation scores to scores.csv\")" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/test/notebook/testfiles/main.py b/test/notebook/testfiles/main.py new file mode 100644 index 00000000..14c9db5e --- /dev/null +++ b/test/notebook/testfiles/main.py @@ -0,0 +1,512 @@ +import os +import sys +import time +import random +import numpy as np +import pandas as pd + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader + +import timm +import albumentations as A +from albumentations.pytorch import ToTensorV2 + +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import roc_auc_score, confusion_matrix + +import cv2 +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument('--debug', action='store_true', help='Run in debug mode') +args = parser.parse_args() +DEBUG = args.debug + +SEED = 2024 +np.random.seed(SEED) +random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +TRAIN_DIR = './workspace_input/train/' +TEST_DIR = './workspace_input/test/' +TRAIN_CSV = './workspace_input/train.csv' +SAMPLE_SUB_PATH = './workspace_input/sample_submission.csv' +MODEL_DIR = 'models/' +os.makedirs(MODEL_DIR, exist_ok=True) + +def print_eda(train_df): + print("=== Start of EDA part ===") + print("Shape of train.csv:", train_df.shape) + print("First 5 rows:\n", train_df.head()) + print("Column data types:\n", train_df.dtypes) + print("Missing values per column:\n", train_df.isnull().sum()) + print("Unique values per column:") + for col in train_df.columns: + print(f" - {col}: {train_df[col].nunique()}") + label_counts = train_df['has_cactus'].value_counts() + print("Label distribution (has_cactus):") + print(label_counts) + pos, neg = label_counts.get(1, 0), label_counts.get(0, 0) + total = pos + neg + if total > 0: + print(f" Positive:Negative ratio: {pos}:{neg} ({pos/total:.3f}:{neg/total:.3f})") + print(f" Percentage positive: {pos/total*100:.2f}%") + else: + print(" No data found.") + print("Image filename examples:", train_df['id'].unique()[:5]) + print("=== End of EDA part ===") + +class CactusDataset(Dataset): + def __init__(self, image_ids, labels=None, id2path=None, transforms=None): + self.image_ids = image_ids + self.labels = labels + self.id2path = id2path + self.transforms = transforms + + def __len__(self): + return len(self.image_ids) + + def __getitem__(self, idx): + img_id = self.image_ids[idx] + img_path = self.id2path[img_id] + image = cv2.imread(img_path) + if image is None: + raise RuntimeError(f"Cannot read image at {img_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + if self.transforms: + augmented = self.transforms(image=image) + image = augmented["image"] + if self.labels is not None: + label = self.labels[idx] + return image, label, img_id + else: + return image, img_id + +def get_transforms(mode='train'): + # Correct Cutout: Albumentations v1.4.15 provides 'Cutout' as a class, but not always in the root. + # Defensive import; fallback to the most robust method for v1.4.15 + imagenet_mean = [0.485, 0.456, 0.406] + imagenet_std = [0.229, 0.224, 0.225] + if mode == 'train': + min_frac, max_frac = 0.05, 0.2 + min_cut = int(300 * min_frac) + max_cut = int(300 * max_frac) + # There is no A.Cutout in v1.4.15 root, but A.augmentations.transforms.Cutout exists. + try: + from albumentations.augmentations.transforms import Cutout + have_cutout = True + except ImportError: + have_cutout = False + this_cut_h = random.randint(min_cut, max_cut) + this_cut_w = random.randint(min_cut, max_cut) + cutout_fill = [int(255 * m) for m in imagenet_mean] + tforms = [ + A.RandomResizedCrop(300, 300, scale=(0.7, 1.0), ratio=(0.8, 1.2), p=1.0), + A.Rotate(limit=30, p=0.8), + ] + if have_cutout: + tforms.append( + Cutout( + num_holes=1, + max_h_size=this_cut_h, + max_w_size=this_cut_w, + fill_value=cutout_fill, # RGB image in albumentations requires [R,G,B] + always_apply=False, + p=0.7 + ) + ) + else: + # No available Cutout, so fallback to no cutout but emit warning + print("WARNING: albumentations.Cutout not found, continuing without Cutout augmentation") + tforms.extend([ + A.RandomContrast(limit=0.2, p=0.5), + A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.1), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + return A.Compose(tforms) + else: + return A.Compose([ + A.Resize(300, 300), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + +def get_dataloader(dataset, batch_size, shuffle=False, num_workers=4, pin_memory=True): + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=pin_memory + ) + +def get_efficientnet_b3(dropout_rate=0.3): + model = timm.create_model('efficientnet_b3', pretrained=True) + n_in = model.classifier.in_features if hasattr(model, "classifier") else model.fc.in_features + model.classifier = nn.Sequential( + nn.Dropout(dropout_rate), + nn.Linear(n_in, 1) + ) + return model + +def compute_class_weight(y): + counts = np.bincount(y) + if len(counts) < 2: + counts = np.pad(counts, (0, 2-len(counts)), constant_values=0) + n_pos, n_neg = counts[1], counts[0] + total = n_pos + n_neg + minority, majority = min(n_pos, n_neg), max(n_pos, n_neg) + ratio = majority / (minority + 1e-10) + need_weights = ratio > 2 + weights = None + if need_weights: + inv_freq = [1 / (n_neg + 1e-10), 1 / (n_pos + 1e-10)] + s = sum(inv_freq) + weights = [w / s * 2 for w in inv_freq] + return weights, n_pos, n_neg, ratio, need_weights + +def train_one_epoch(model, loss_fn, optimizer, scheduler, dataloader, device, class_weights): + model.train() + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + if scheduler is not None: + scheduler.step() + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + avg_loss = total_loss / total_samples + return avg_loss + +@torch.no_grad() +def eval_model(model, loss_fn, dataloader, device, class_weights): + model.eval() + y_true, y_pred = [], [] + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + probs = torch.sigmoid(logits) + y_true.append(labels.cpu().numpy()) + y_pred.append(probs.cpu().numpy()) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + y_true = np.vstack(y_true).reshape(-1) + y_pred = np.vstack(y_pred).reshape(-1) + avg_loss = total_loss / total_samples + return avg_loss, y_true, y_pred + +def confusion_info(y_true, y_pred, threshold=0.5): + preds = (y_pred > threshold).astype(int) + cm = confusion_matrix(y_true, preds) + return cm + +def inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, class_weights, need_weights, + BATCH_SIZE, N_WORKERS, cv_fold): + oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] + for fold in range(cv_fold): + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + val_ds = CactusDataset(val_img_ids, val_labels, id2path=train_id2path, transforms=get_transforms("val")) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(val_auc) + print(f"Reloaded fold {fold}, OOF Validation AUC={val_auc:.5f}") + + all_oof_true = np.concatenate(oof_true) + all_oof_pred = np.concatenate(oof_pred) + oof_auc = roc_auc_score(all_oof_true, all_oof_pred) + oof_cm = confusion_info(all_oof_true, all_oof_pred) + print(f"OOF ROC-AUC (from loaded models): {oof_auc:.5f}") + print(f"OOF Confusion Matrix:\n{oof_cm}") + + test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") + ) + test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + test_pred_list = [] + for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") + test_probs = np.mean(test_pred_list, axis=0) + + submission = pd.read_csv(SAMPLE_SUB_PATH) + submission['has_cactus'] = test_probs + submission.to_csv('submission.csv', index=False) + print(f"Saved submission.csv in required format with {len(submission)} rows.") + + scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] + }) + scores_df.set_index('Model', inplace=True) + scores_df.to_csv("scores.csv") + print(f"Saved cross-validation scores to scores.csv") + +def main(): + print("Section: Data Loading and Preprocessing") + # This section loads the train and test data, performs EDA, and prepares the dataset. + try: + train_df = pd.read_csv(TRAIN_CSV) + except Exception as e: + print(f"Failed to load train.csv: {e}") + sys.exit(1) + print_eda(train_df) + + train_id2path = {img_id: os.path.join(TRAIN_DIR, img_id) for img_id in train_df['id']} + try: + sample_sub = pd.read_csv(SAMPLE_SUB_PATH) + except Exception as e: + print(f"Failed to load sample_submission.csv: {e}") + sys.exit(1) + test_img_ids = list(sample_sub['id']) + test_id2path = {img_id: os.path.join(TEST_DIR, img_id) for img_id in test_img_ids} + print(f"Loaded {len(train_id2path)} train images, {len(test_id2path)} test images.") + + y_train = train_df['has_cactus'].values + class_weights, n_pos, n_neg, imbalance_ratio, need_weights = compute_class_weight(y_train) + print(f"Class stats: Pos={n_pos}, Neg={n_neg}, Imbalance Ratio(majority/minority)={imbalance_ratio:.3f}") + print(f"Use class weights: {need_weights}, Class weights: {class_weights if class_weights is not None else '[1.0,1.0]'}") + if class_weights is not None: + np.save(os.path.join(MODEL_DIR, "class_weights.npy"), class_weights) + + print("Section: Feature Engineering") + train_df = train_df.copy() + cv_fold = 5 + skf = StratifiedKFold(n_splits=cv_fold, shuffle=True, random_state=SEED) + folds = np.zeros(len(train_df), dtype=np.int32) + for idx, (_, val_idx) in enumerate(skf.split(train_df['id'], train_df['has_cactus'])): + folds[val_idx] = idx + train_df['fold'] = folds + print(f"Assigned stratified {cv_fold}-fold indices. Fold sample counts:") + for f in range(cv_fold): + dist = train_df.loc[train_df['fold'] == f, 'has_cactus'].value_counts().to_dict() + print(f" Fold {f}: n={len(train_df[train_df['fold'] == f])} class dist={dist}") + + print("Section: Model Training and Evaluation") + dropout_rate = round(random.uniform(0.2, 0.5), 2) + print(f"Model config: EfficientNet-B3, Image size 300, Head dropout={dropout_rate}") + + if DEBUG: + print("DEBUG mode: using 10% subsample and 1 epoch (per fold)") + sample_frac = 0.10 + sampled_idxs = [] + for f in range(cv_fold): + fold_idx = train_df.index[train_df['fold'] == f].tolist() + fold_labels = train_df.loc[fold_idx, 'has_cactus'].values + idx_pos = [i for i, l in zip(fold_idx, fold_labels) if l == 1] + idx_neg = [i for i, l in zip(fold_idx, fold_labels) if l == 0] + n_pos = max(1, int(sample_frac * len(idx_pos))) + n_neg = max(1, int(sample_frac * len(idx_neg))) + if len(idx_pos) > 0: + sampled_idxs += np.random.choice(idx_pos, n_pos, replace=False).tolist() + if len(idx_neg) > 0: + sampled_idxs += np.random.choice(idx_neg, n_neg, replace=False).tolist() + train_df = train_df.loc[sampled_idxs].reset_index(drop=True) + print(f"DEBUG subsample shape: {train_df.shape}") + debug_epochs = 1 + else: + debug_epochs = None + + BATCH_SIZE = 64 if torch.cuda.is_available() else 32 + N_WORKERS = 4 if torch.cuda.is_available() else 1 + EPOCHS = 20 if not DEBUG else debug_epochs + MIN_EPOCHS = 5 if not DEBUG else 1 + EARLY_STOP_PATIENCE = 7 if not DEBUG else 2 + LR = 1e-3 + + model_files = [os.path.join(MODEL_DIR, f"efficientnet_b3_fold{f}.pt") for f in range(cv_fold)] + if all([os.path.exists(f) for f in model_files]): + print("All fold models found in models/. Running inference and file saving only (no retrain).") + inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, + class_weights, need_weights, BATCH_SIZE, N_WORKERS, cv_fold) + return + + oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] + start_time = time.time() if DEBUG else None + + for fold in range(cv_fold): + print(f"\n=== FOLD {fold} TRAINING ===") + df_train = train_df[train_df['fold'] != fold].reset_index(drop=True) + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + print(f"Train size: {df_train.shape[0]}, Val size: {df_val.shape[0]}") + train_img_ids = df_train['id'].tolist() + train_labels = df_train['has_cactus'].values + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + + train_ds = CactusDataset( + train_img_ids, train_labels, + id2path=train_id2path, + transforms=get_transforms("train") + ) + val_ds = CactusDataset( + val_img_ids, val_labels, + id2path=train_id2path, + transforms=get_transforms("val") + ) + train_loader = get_dataloader(train_ds, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + optimizer = optim.AdamW(model.parameters(), lr=LR) + scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + best_auc = -np.inf + best_epoch = -1 + best_model_state = None + patience = 0 + + for epoch in range(EPOCHS): + train_loss = train_one_epoch( + model, loss_fn, optimizer, scheduler, train_loader, DEVICE, fold_class_weights) + val_loss, val_true, val_pred = eval_model( + model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + cm = confusion_info(val_true, val_pred) + print(f"Epoch {epoch+1:02d}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_auc={val_auc:.4f}") + print(f" Val confusion_matrix (rows:true [0,1]; cols:pred [0,1]):\n{cm}") + if val_auc > best_auc: + best_auc = val_auc + best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} + best_epoch = epoch + patience = 0 + else: + patience += 1 + if DEBUG and epoch + 1 >= debug_epochs: + break + if (epoch + 1) >= MIN_EPOCHS and patience >= EARLY_STOP_PATIENCE: + print(f"Early stopping at epoch {epoch+1}, best_epoch={best_epoch+1}.") + break + + model.load_state_dict(best_model_state) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + torch.save(model.state_dict(), fold_model_path) + print(f"Saved best model for fold {fold} at {fold_model_path} (best_auc={best_auc:.5f}, best_epoch={best_epoch+1})") + + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(best_auc) + print(f"OOF stored for fold {fold}, Validation AUC={best_auc:.5f}") + + end_time = time.time() if DEBUG else None + if DEBUG: + debug_time = end_time - start_time + estimated_time = (1 / 0.1) * (EPOCHS / debug_epochs) * debug_time + print("=== Start of Debug Information ===") + print(f"debug_time: {debug_time:.1f}") + print(f"estimated_time: {estimated_time:.1f}") + print("=== End of Debug Information ===") + + print("Section: Ensemble Strategy and Final Predictions") + all_oof_true = np.concatenate(oof_true) + all_oof_pred = np.concatenate(oof_pred) + oof_auc = roc_auc_score(all_oof_true, all_oof_pred) + oof_cm = confusion_info(all_oof_true, all_oof_pred) + print(f"OOF ROC-AUC: {oof_auc:.5f}") + print(f"OOF Confusion Matrix:\n{oof_cm}") + + test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") + ) + test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + test_pred_list = [] + for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") + test_probs = np.mean(test_pred_list, axis=0) + + print("Section: Submission File Generation") + submission = pd.read_csv(SAMPLE_SUB_PATH) + submission['has_cactus'] = test_probs + submission.to_csv('submission.csv', index=False) + print(f"Saved submission.csv in required format with {len(submission)} rows.") + + scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] + }) + scores_df.set_index('Model', inplace=True) + scores_df.to_csv("scores.csv") + print(f"Saved cross-validation scores to scores.csv") + +if __name__ == "__main__": + main() diff --git a/test/notebook/testfiles/main2.ipynb b/test/notebook/testfiles/main2.ipynb new file mode 100644 index 00000000..0b13b572 --- /dev/null +++ b/test/notebook/testfiles/main2.ipynb @@ -0,0 +1,609 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "3314929a", + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "# hack to allow argparse to work in notebook\n", + "sys.argv = [\"main.py\"]\n", + "\n", + "import os\n", + "import time\n", + "import argparse\n", + "import random\n", + "import numpy as np\n", + "import pandas as pd\n", + "from PIL import Image\n", + "from glob import glob\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from torch.utils.data import Dataset, DataLoader\n", + "import torchvision\n", + "\n", + "import albumentations as A\n", + "from albumentations.pytorch import ToTensorV2\n", + "import cv2\n", + "\n", + "from sklearn.model_selection import StratifiedShuffleSplit\n", + "from sklearn.metrics import log_loss\n", + "\n", + "# ========= Debug mode handling ==========\n", + "parser = argparse.ArgumentParser()\n", + "parser.add_argument('--debug', action='store_true', help='Run in debug mode')\n", + "args = parser.parse_args()\n", + "DEBUG = False\n", + "if args.debug:\n", + " DEBUG = True\n", + "\n", + "# ========= Set random seed for reproducibility ==========\n", + "def seed_everything(seed=42):\n", + " random.seed(seed)\n", + " np.random.seed(seed)\n", + " torch.manual_seed(seed)\n", + " torch.cuda.manual_seed_all(seed)\n", + "seed_everything(42)\n", + "\n", + "DATA_DIR = './workspace_input/'\n", + "TRAIN_CSV = os.path.join(DATA_DIR, 'train.csv')\n", + "TRAIN_DIR = os.path.join(DATA_DIR, 'train/')\n", + "TEST_DIR = os.path.join(DATA_DIR, 'test/')\n", + "SAMPLE_SUB_CSV = os.path.join(DATA_DIR, 'sample_submission.csv')\n", + "MODEL_DIR = 'models/'\n", + "SUBMISSION_PATH = 'submission.csv'\n", + "SCORES_PATH = 'scores.csv'\n", + "\n", + "if not os.path.exists(MODEL_DIR):\n", + " os.makedirs(MODEL_DIR, exist_ok=True)\n" + ] + }, + { + "cell_type": "markdown", + "id": "2e42af1b", + "metadata": {}, + "source": [ + "## Data Loading and Preprocessing\n", + "Load train.csv and list image files in train/ and test/\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa7c7a55", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Data Loading and Preprocessing\")\n", + "try:\n", + " train_df = pd.read_csv(TRAIN_CSV)\n", + "except Exception as e:\n", + " print(f\"Error loading train.csv: {e}\")\n", + " exit(1)\n", + "\n", + "try:\n", + " train_image_files = set(os.listdir(TRAIN_DIR))\n", + "except Exception as e:\n", + " print(f\"Error listing train dir: {e}\")\n", + " exit(1)\n", + "\n", + "try:\n", + " test_image_files = set(os.listdir(TEST_DIR))\n", + "except Exception as e:\n", + " print(f\"Error listing test dir: {e}\")\n", + " exit(1)\n", + "\n", + "# Confirm train_df ids and image files match\n", + "train_df = train_df[train_df['id'].isin(train_image_files)].reset_index(drop=True)\n", + "test_image_files = sorted(list(test_image_files))\n", + "\n", + "try:\n", + " sample_submission = pd.read_csv(SAMPLE_SUB_CSV)\n", + " SUB_COLS = sample_submission.columns.tolist()\n", + "except Exception as e:\n", + " print(f\"Error reading sample_submission.csv: {e}\")\n", + " SUB_COLS = ['id', 'has_cactus']" + ] + }, + { + "cell_type": "markdown", + "id": "450bb94b", + "metadata": {}, + "source": [ + "## Exploratory Data Analysis (EDA)\n", + "EDA Output Generation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ea29a876", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Exploratory Data Analysis (EDA)\")\n", + "n_train = len(train_df)\n", + "n_test = len(test_image_files)\n", + "train_ids = train_df['id'].tolist()\n", + "eda_content = []\n", + "eda_content.append(\"=== Start of EDA part ===\")\n", + "eda_content.append(f\"Train.csv shape: {train_df.shape}\")\n", + "eda_content.append(f\"First 5 rows:\\n{train_df.head(5).to_string(index=False)}\")\n", + "eda_content.append(f\"\\nData types:\\n{train_df.dtypes.to_string()}\")\n", + "eda_content.append(f\"\\nMissing values:\\n{train_df.isnull().sum().to_string()}\")\n", + "eda_content.append(f\"\\nUnique values per column:\\n{train_df.nunique()}\")\n", + "class_dist = train_df['has_cactus'].value_counts().sort_index()\n", + "eda_content.append(f\"\\nTarget distribution:\\n{class_dist.to_string()}\")\n", + "eda_content.append(f\"\\nBalance ratio (majority/minority): {class_dist.max()/class_dist.min():.2f}\")\n", + "eda_content.append(f\"\\nTotal train images in 'train/' folder: {len(train_image_files)}\")\n", + "eda_content.append(f\"Total test images in 'test/' folder: {len(test_image_files)}\")\n", + "eda_content.append(f\"All train.csv ids found in train/: {all(i in train_image_files for i in train_df['id'])}\")\n", + "eda_content.append(f\"Sample of train image filename: {train_df['id'].iloc[0]}\")\n", + "eda_content.append(f\"Sample of test image filename: {test_image_files[0]}\")\n", + "eda_content.append(\"Image format: assumed all JPG, size like 32x32 px (EfficientNet expects resize to 224x224)\")\n", + "eda_content.append(\"No missing values detected in train.csv; binary target (0=no cactus, 1=has cactus).\")\n", + "eda_content.append(\"No duplicates in train.csv ids. Appears to be balanced.\")\n", + "eda_content.append(\"=== End of EDA part ===\")\n", + "print('\\n'.join(eda_content))" + ] + }, + { + "cell_type": "markdown", + "id": "6723009f", + "metadata": {}, + "source": [ + "## Feature Engineering - Green Mask Channel\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8e24b0ca", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Feature Engineering - Green Mask Channel\")\n", + "def green_mask(img_bgr):\n", + " hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)\n", + " lower = np.array([35, 51, 41], dtype=np.uint8)\n", + " upper = np.array([85, 255, 255], dtype=np.uint8)\n", + " mask = cv2.inRange(hsv, lower, upper)\n", + " mask = (mask > 0).astype(np.uint8)\n", + " return mask[..., None]\n", + "\n", + "def load_img_as_numpy_with_mask(filepath):\n", + " try:\n", + " img_bgr = cv2.imread(filepath, cv2.IMREAD_COLOR)\n", + " if img_bgr is None:\n", + " raise ValueError(f\"cv2.imread failed for {filepath}\")\n", + " img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n", + " mask = green_mask(img_bgr)\n", + " img4 = np.concatenate([img_rgb, mask*255], axis=2)\n", + " return img4\n", + " except Exception as e:\n", + " print(f\"Error reading {filepath}: {e}\")\n", + " return np.zeros((32, 32, 4), dtype=np.uint8)\n", + "\n", + "test_ids = test_image_files" + ] + }, + { + "cell_type": "markdown", + "id": "9345e92a", + "metadata": {}, + "source": [ + "## Data Augmentation and Transform Pipeline\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f051fe0e", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Data Augmentation and Transform Pipeline\")\n", + "\n", + "IMG_SIZE = 224\n", + "MEAN = [0.485, 0.456, 0.406, 0.0]\n", + "STD = [0.229, 0.224, 0.225, 1.0]\n", + "\n", + "def get_transforms(mode='train'):\n", + " if mode == 'train':\n", + " aug = [\n", + " A.Resize(IMG_SIZE, IMG_SIZE),\n", + " A.OneOf([\n", + " A.Affine(rotate=(-25,25), shear={'x':(-8,8),'y':(-8,8)}, scale=(0.9,1.1), translate_percent={\"x\":(-0.1,0.1),\"y\":(-0.1,0.1)}),\n", + " A.NoOp()],\n", + " p=0.5\n", + " ),\n", + " A.HorizontalFlip(p=0.5),\n", + " A.VerticalFlip(p=0.5),\n", + " A.RandomBrightnessContrast(brightness_limit=0.18, contrast_limit=0.15, p=0.5),\n", + " A.HueSaturationValue(hue_shift_limit=7, sat_shift_limit=15, val_shift_limit=10, p=0.5),\n", + " A.GaussianNoise(var_limit=(10.0, 30.0), p=0.5),\n", + " A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.),\n", + " ToTensorV2(transpose_mask=True),\n", + " ]\n", + " return A.Compose(aug)\n", + " else:\n", + " aug = [\n", + " A.Resize(IMG_SIZE, IMG_SIZE),\n", + " A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.),\n", + " ToTensorV2(transpose_mask=True),\n", + " ]\n", + " return A.Compose(aug)" + ] + }, + { + "cell_type": "markdown", + "id": "0d67fb3a", + "metadata": {}, + "source": [ + "## Dataset and DataLoader Construction\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18bbcedb", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Dataset and DataLoader Construction\")\n", + "\n", + "class CactusDataset(Dataset):\n", + " def __init__(self, img_ids, img_dir, labels=None, transform=None, cache=False):\n", + " self.img_ids = img_ids\n", + " self.img_dir = img_dir\n", + " self.labels = labels # None for test\n", + " self.transform = transform\n", + " self.cache = cache\n", + " self._cache = {}\n", + " def __len__(self):\n", + " return len(self.img_ids)\n", + " def __getitem__(self, idx):\n", + " img_id = self.img_ids[idx]\n", + " if self.cache and img_id in self._cache:\n", + " img4 = self._cache[img_id]\n", + " else:\n", + " img_path = os.path.join(self.img_dir, img_id)\n", + " img4 = load_img_as_numpy_with_mask(img_path)\n", + " if self.cache:\n", + " self._cache[img_id] = img4\n", + " transformed = self.transform(image=img4)\n", + " img = transformed['image']\n", + " if self.labels is not None:\n", + " label = float(self.labels[idx])\n", + " return img, label\n", + " else:\n", + " return img, img_id\n", + "\n", + "split_seed = 42\n", + "splitter = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=split_seed)\n", + "try:\n", + " split = next(splitter.split(train_df['id'], train_df['has_cactus']))\n", + " tr_indices, val_indices = split\n", + "except Exception as e:\n", + " print(f'Stratified split failed ({e}), falling back to random split')\n", + " indices = np.arange(len(train_df))\n", + " np.random.shuffle(indices)\n", + " n_val = int(0.2 * len(train_df))\n", + " val_indices = indices[:n_val]\n", + " tr_indices = indices[n_val:]\n", + "\n", + "# Sampling, only in debug mode: sample *after* split\n", + "if DEBUG:\n", + " tr_sample_size = max(2, int(0.1 * len(tr_indices)))\n", + " val_sample_size = max(2, int(0.1 * len(val_indices)))\n", + " tr_indices = np.random.choice(tr_indices, tr_sample_size, replace=False)\n", + " val_indices = np.random.choice(val_indices, val_sample_size, replace=False)\n", + "\n", + "tr_ids = train_df.iloc[tr_indices]['id'].tolist()\n", + "val_ids = train_df.iloc[val_indices]['id'].tolist()\n", + "tr_lbls = train_df.iloc[tr_indices]['has_cactus'].tolist()\n", + "val_lbls = train_df.iloc[val_indices]['has_cactus'].tolist()\n", + "\n", + "# For reproducibility and fast debug, cache only in debug for train/val.\n", + "train_ds = CactusDataset(tr_ids, TRAIN_DIR, tr_lbls, transform=get_transforms('train'), cache=(DEBUG))\n", + "val_ds = CactusDataset(val_ids, TRAIN_DIR, val_lbls, transform=get_transforms('val'), cache=(DEBUG))\n", + "test_ds = CactusDataset(test_ids, TEST_DIR, labels=None, transform=get_transforms('val'), cache=False)\n", + "\n", + "BATCH_SIZE = 32 if not DEBUG else 8\n", + "NUM_WORKERS = min(4, os.cpu_count())\n", + "\n", + "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, drop_last=False, num_workers=NUM_WORKERS, pin_memory=True)\n", + "val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, drop_last=False, num_workers=NUM_WORKERS, pin_memory=True)\n", + "test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE*2, shuffle=False, drop_last=False, num_workers=NUM_WORKERS, pin_memory=True)" + ] + }, + { + "cell_type": "markdown", + "id": "5f5b5efd", + "metadata": {}, + "source": [ + "## Model Definition and Adaptation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be8a39fa", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Model Definition and Adaptation\")\n", + "class EfficientNetB0_4ch(nn.Module):\n", + " def __init__(self, pretrained=True):\n", + " super().__init__()\n", + " from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights\n", + " if pretrained:\n", + " wts = EfficientNet_B0_Weights.DEFAULT\n", + " net = efficientnet_b0(weights=wts)\n", + " else:\n", + " net = efficientnet_b0(weights=None)\n", + " old_conv = net.features[0][0]\n", + " new_conv = nn.Conv2d(4, old_conv.out_channels, kernel_size=old_conv.kernel_size,\n", + " stride=old_conv.stride, padding=old_conv.padding, bias=False)\n", + " with torch.no_grad():\n", + " new_conv.weight[:, :3] = old_conv.weight\n", + " mean_wt = torch.mean(old_conv.weight, dim=1, keepdim=True)\n", + " new_conv.weight[:, 3:4] = mean_wt\n", + " net.features[0][0] = new_conv\n", + " self.features = net.features\n", + " self.avgpool = net.avgpool\n", + " inner_dim = net.classifier[1].in_features\n", + " self.head = nn.Sequential(\n", + " nn.Dropout(0.3),\n", + " nn.Linear(inner_dim, 1)\n", + " )\n", + " def forward(self, x):\n", + " x = self.features(x)\n", + " x = self.avgpool(x)\n", + " x = torch.flatten(x, 1)\n", + " x = self.head(x)\n", + " return x\n", + "\n", + "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", + "MODEL_TRAINED_FILE = os.path.join(MODEL_DIR, 'efficientnet_b0_best.pth')\n", + "scaler = torch.cuda.amp.GradScaler() if torch.cuda.is_available() else None\n", + "\n", + "# Timing stats for debug regardless path\n", + "debug_time = None\n", + "estimated_time = None\n", + "\n", + "NEED_TRAIN = not (os.path.isfile(MODEL_TRAINED_FILE))\n", + "if not NEED_TRAIN:\n", + " print(\"Model checkpoint detected, will use it for inference!\")\n", + " model = EfficientNetB0_4ch(pretrained=False).to(device)\n", + " state = torch.load(MODEL_TRAINED_FILE, map_location=device)\n", + " model.load_state_dict(state['model'])\n", + " # If in debug, set fake small debug_time for inference-only, as required for compliance.\n", + " if DEBUG:\n", + " debug_time = 1.0\n", + " scale = (1/0.1) * (1 if DEBUG else 20)\n", + " estimated_time = debug_time * scale\n", + "else:\n", + " print(\"Model checkpoint not found, proceeding to training...\")\n", + " print(\"Section: Training: Staged Fine-Tuning with Discriminative LRs\")\n", + " model = EfficientNetB0_4ch(pretrained=True).to(device)\n", + " criterion = nn.BCEWithLogitsLoss()\n", + " backbone_params = []\n", + " mid_params = []\n", + " head_params = list(model.head.parameters())\n", + " for i, m in enumerate(model.features):\n", + " if i <= 2:\n", + " backbone_params += list(m.parameters())\n", + " elif 3 <= i <= 5:\n", + " mid_params += list(m.parameters())\n", + " def set_requires_grad(modules, req):\n", + " for m in modules:\n", + " for param in m.parameters():\n", + " param.requires_grad = req\n", + " set_requires_grad([model.features], False)\n", + " set_requires_grad([model.head], True)\n", + " EPOCHS = 20 if not DEBUG else 1\n", + " patience = 5\n", + " optimizer = optim.Adam(model.head.parameters(), lr=5e-4, weight_decay=1e-5)\n", + " scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)\n", + " best_loss = float('inf')\n", + " best_state = None\n", + " patience_counter = 0\n", + " start_time = time.time() if DEBUG else None\n", + " for epoch in range(EPOCHS):\n", + " print(f\"Epoch {epoch+1}/{EPOCHS}\")\n", + " if epoch == 3:\n", + " set_requires_grad([model.features[3], model.features[4], model.features[5]], True)\n", + " optimizer = optim.Adam([\n", + " {'params': backbone_params, 'lr': 1e-4},\n", + " {'params': mid_params, 'lr': 2e-4},\n", + " {'params': head_params, 'lr':5e-4},\n", + " ], weight_decay=1e-5)\n", + " scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS-epoch)\n", + " print(\"Unfroze mid layers of EfficientNet for fine-tuning.\")\n", + " elif epoch == 6:\n", + " set_requires_grad([model.features], True)\n", + " print(\"Unfroze all layers of EfficientNet for full fine-tuning.\")\n", + "\n", + " model.train()\n", + " tr_loss = 0.\n", + " tr_cnt = 0\n", + " for imgs, lbls in train_loader:\n", + " imgs = imgs.to(device)\n", + " lbls = lbls.to(device).view(-1,1)\n", + " optimizer.zero_grad()\n", + " if scaler is not None:\n", + " with torch.cuda.amp.autocast():\n", + " outs = model(imgs)\n", + " loss = criterion(outs, lbls)\n", + " scaler.scale(loss).backward()\n", + " scaler.step(optimizer)\n", + " scaler.update()\n", + " else:\n", + " outs = model(imgs)\n", + " loss = criterion(outs, lbls)\n", + " loss.backward()\n", + " optimizer.step()\n", + " tr_loss += loss.item() * imgs.size(0)\n", + " tr_cnt += imgs.size(0)\n", + " if scheduler is not None:\n", + " scheduler.step()\n", + "\n", + " tr_loss = tr_loss / tr_cnt\n", + "\n", + " model.eval()\n", + " val_loss = 0.\n", + " val_cnt = 0\n", + " all_val_lbls = []\n", + " all_val_preds = []\n", + " with torch.no_grad():\n", + " for imgs, lbls in val_loader:\n", + " imgs = imgs.to(device)\n", + " lbls = lbls.cpu().numpy()\n", + " outs = model(imgs).cpu().squeeze().numpy()\n", + " preds = 1/(1 + np.exp(-outs))\n", + " loss = criterion(torch.tensor(outs).view(-1,1), torch.tensor(lbls).view(-1,1)).item()\n", + " val_loss += loss * imgs.size(0)\n", + " val_cnt += imgs.size(0)\n", + " all_val_lbls.append(lbls)\n", + " all_val_preds.append(preds)\n", + " val_loss = val_loss / val_cnt\n", + " all_val_lbls = np.concatenate(all_val_lbls)\n", + " all_val_preds = np.concatenate(all_val_preds)\n", + " try:\n", + " val_logloss = log_loss(all_val_lbls, all_val_preds, eps=1e-7)\n", + " except Exception as ex:\n", + " val_logloss = float('inf')\n", + " print(\"Error computing log_loss on val:\", ex)\n", + "\n", + " print(f\"Train Loss: {tr_loss:.5f} | Val Loss (BCE): {val_loss:.5f} | Val LogLoss: {val_logloss:.5f}\")\n", + "\n", + " if val_logloss < best_loss:\n", + " best_loss = val_logloss\n", + " best_state = {\n", + " 'model': model.state_dict(),\n", + " 'epoch': epoch,\n", + " 'val_loss': best_loss,\n", + " }\n", + " torch.save(best_state, MODEL_TRAINED_FILE)\n", + " patience_counter = 0\n", + " print(f\"Best model saved. (epoch {epoch+1}, val_logloss={val_logloss:.5f})\")\n", + " else:\n", + " patience_counter += 1\n", + " print(f\"No improvement. Early stopping patience: {patience_counter}/{patience}\")\n", + "\n", + " if patience_counter >= patience:\n", + " print(f\"Early stopping triggered at epoch {epoch+1}.\")\n", + " break\n", + " if DEBUG and start_time is not None:\n", + " end_time = time.time()\n", + " debug_time = end_time - start_time\n", + " # Compute estimated time: (fractional data)*(epochs) compared\n", + " sample_factor = 0.1\n", + " scale = (1/sample_factor) * (20 if not DEBUG else 1)\n", + " estimated_time = debug_time * scale\n", + " # Reload best model for evaluation\n", + " state = torch.load(MODEL_TRAINED_FILE, map_location=device)\n", + " model.load_state_dict(state['model'])" + ] + }, + { + "cell_type": "markdown", + "id": "0d98a34c", + "metadata": {}, + "source": [ + "## Validation Evaluation and Metric Calculation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6b2bfe97", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Validation Evaluation and Metric Calculation\")\n", + "model.eval()\n", + "val_lbls, val_prs = [], []\n", + "with torch.no_grad():\n", + " for imgs, lbls in val_loader:\n", + " imgs = imgs.to(device)\n", + " outs = model(imgs).cpu().squeeze().numpy()\n", + " prs = 1/(1+np.exp(-outs))\n", + " val_lbls.append(lbls.numpy())\n", + " val_prs.append(prs)\n", + "val_lbls = np.concatenate(val_lbls)\n", + "val_prs = np.concatenate(val_prs)\n", + "try:\n", + " val_logloss = log_loss(val_lbls, val_prs, eps=1e-7)\n", + "except Exception as ex:\n", + " val_logloss = float('inf')\n", + " print(\"Error computing log_loss on validation:\", ex)\n", + "print(f\"Final best model log loss on validation split: {val_logloss:.6f}\")\n", + "scores = pd.DataFrame(\n", + " {'Model': ['efficientnet_b0', 'ensemble'], 'LogLoss': [val_logloss, val_logloss]}\n", + ").set_index('Model')\n", + "scores.to_csv(SCORES_PATH)\n", + "print(f\"Saved scores.csv with validation log loss.\")" + ] + }, + { + "cell_type": "markdown", + "id": "6a45e9cb", + "metadata": {}, + "source": [ + "## Prediction and Submission Generation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6bc7e8e0", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Section: Prediction and Submission Generation\")\n", + "model.eval()\n", + "test_probs = []\n", + "test_ids_ordered = []\n", + "with torch.no_grad():\n", + " for imgs, img_ids in test_loader:\n", + " imgs = imgs.to(device)\n", + " outs = model(imgs).cpu().squeeze().numpy()\n", + " prs = 1/(1+np.exp(-outs))\n", + " if isinstance(img_ids, list) or isinstance(img_ids, np.ndarray):\n", + " test_ids_ordered += list(img_ids)\n", + " else:\n", + " test_ids_ordered.append(img_ids)\n", + " test_probs.extend(np.array(prs).ravel().tolist())\n", + "submit_df = pd.DataFrame({'id': test_ids_ordered, 'has_cactus': test_probs})\n", + "submit_df = submit_df.set_index('id')\n", + "try:\n", + " submit_df = submit_df.reindex(sample_submission['id']).reset_index()\n", + "except Exception:\n", + " submit_df = submit_df.reset_index()\n", + "submit_df['has_cactus'] = submit_df['has_cactus'].clip(0,1)\n", + "submit_df.to_csv(SUBMISSION_PATH, index=False, float_format='%.6f')\n", + "print(f\"Saved submission.csv with {len(submit_df)} rows. Format: {submit_df.columns.tolist()}\")\n", + "\n", + "# === Debug info output, always print in debug mode, even if only inference ===\n", + "if DEBUG:\n", + " if debug_time is None:\n", + " debug_time = 1.0\n", + " scale = (1/0.1)*(1 if DEBUG else 20)\n", + " estimated_time = debug_time * scale\n", + " print(\"=== Start of Debug Information ===\")\n", + " print(f\"debug_time: {debug_time}\")\n", + " print(f\"estimated_time: {estimated_time}\")\n", + " print(\"=== End of Debug Information ===\")" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/test/notebook/testfiles/main2.py b/test/notebook/testfiles/main2.py new file mode 100644 index 00000000..5ce74205 --- /dev/null +++ b/test/notebook/testfiles/main2.py @@ -0,0 +1,466 @@ +import os +import time +import argparse +import random +import numpy as np +import pandas as pd +from PIL import Image +from glob import glob + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader +import torchvision + +import albumentations as A +from albumentations.pytorch import ToTensorV2 +import cv2 + +from sklearn.model_selection import StratifiedShuffleSplit +from sklearn.metrics import log_loss + +# ========= Debug mode handling ========== +parser = argparse.ArgumentParser() +parser.add_argument('--debug', action='store_true', help='Run in debug mode') +args = parser.parse_args() +DEBUG = False +if args.debug: + DEBUG = True + +# ========= Set random seed for reproducibility ========== +def seed_everything(seed=42): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) +seed_everything(42) + +def main(): + # ========= Paths ========== + DATA_DIR = './workspace_input/' + TRAIN_CSV = os.path.join(DATA_DIR, 'train.csv') + TRAIN_DIR = os.path.join(DATA_DIR, 'train/') + TEST_DIR = os.path.join(DATA_DIR, 'test/') + SAMPLE_SUB_CSV = os.path.join(DATA_DIR, 'sample_submission.csv') + MODEL_DIR = 'models/' + SUBMISSION_PATH = 'submission.csv' + SCORES_PATH = 'scores.csv' + + if not os.path.exists(MODEL_DIR): + os.makedirs(MODEL_DIR, exist_ok=True) + + print("Section: Data Loading and Preprocessing") + # Load train.csv and list image files in train/ and test/ + try: + train_df = pd.read_csv(TRAIN_CSV) + except Exception as e: + print(f"Error loading train.csv: {e}") + exit(1) + + try: + train_image_files = set(os.listdir(TRAIN_DIR)) + except Exception as e: + print(f"Error listing train dir: {e}") + exit(1) + + try: + test_image_files = set(os.listdir(TEST_DIR)) + except Exception as e: + print(f"Error listing test dir: {e}") + exit(1) + + # Confirm train_df ids and image files match + train_df = train_df[train_df['id'].isin(train_image_files)].reset_index(drop=True) + test_image_files = sorted(list(test_image_files)) + + try: + sample_submission = pd.read_csv(SAMPLE_SUB_CSV) + SUB_COLS = sample_submission.columns.tolist() + except Exception as e: + print(f"Error reading sample_submission.csv: {e}") + SUB_COLS = ['id', 'has_cactus'] + + print("Section: Exploratory Data Analysis (EDA)") + # EDA Output Generation + n_train = len(train_df) + n_test = len(test_image_files) + train_ids = train_df['id'].tolist() + eda_content = [] + eda_content.append("=== Start of EDA part ===") + eda_content.append(f"Train.csv shape: {train_df.shape}") + eda_content.append(f"First 5 rows:\n{train_df.head(5).to_string(index=False)}") + eda_content.append(f"\nData types:\n{train_df.dtypes.to_string()}") + eda_content.append(f"\nMissing values:\n{train_df.isnull().sum().to_string()}") + eda_content.append(f"\nUnique values per column:\n{train_df.nunique()}") + class_dist = train_df['has_cactus'].value_counts().sort_index() + eda_content.append(f"\nTarget distribution:\n{class_dist.to_string()}") + eda_content.append(f"\nBalance ratio (majority/minority): {class_dist.max()/class_dist.min():.2f}") + eda_content.append(f"\nTotal train images in 'train/' folder: {len(train_image_files)}") + eda_content.append(f"Total test images in 'test/' folder: {len(test_image_files)}") + eda_content.append(f"All train.csv ids found in train/: {all(i in train_image_files for i in train_df['id'])}") + eda_content.append(f"Sample of train image filename: {train_df['id'].iloc[0]}") + eda_content.append(f"Sample of test image filename: {test_image_files[0]}") + eda_content.append("Image format: assumed all JPG, size like 32x32 px (EfficientNet expects resize to 224x224)") + eda_content.append("No missing values detected in train.csv; binary target (0=no cactus, 1=has cactus).") + eda_content.append("No duplicates in train.csv ids. Appears to be balanced.") + eda_content.append("=== End of EDA part ===") + print('\n'.join(eda_content)) + + print("Section: Feature Engineering - Green Mask Channel") + def green_mask(img_bgr): + hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) + lower = np.array([35, 51, 41], dtype=np.uint8) + upper = np.array([85, 255, 255], dtype=np.uint8) + mask = cv2.inRange(hsv, lower, upper) + mask = (mask > 0).astype(np.uint8) + return mask[..., None] + + def load_img_as_numpy_with_mask(filepath): + try: + img_bgr = cv2.imread(filepath, cv2.IMREAD_COLOR) + if img_bgr is None: + raise ValueError(f"cv2.imread failed for {filepath}") + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + mask = green_mask(img_bgr) + img4 = np.concatenate([img_rgb, mask*255], axis=2) + return img4 + except Exception as e: + print(f"Error reading {filepath}: {e}") + return np.zeros((32, 32, 4), dtype=np.uint8) + + test_ids = test_image_files + + print("Section: Data Augmentation and Transform Pipeline") + + IMG_SIZE = 224 + MEAN = [0.485, 0.456, 0.406, 0.0] + STD = [0.229, 0.224, 0.225, 1.0] + + def get_transforms(mode='train'): + if mode == 'train': + aug = [ + A.Resize(IMG_SIZE, IMG_SIZE), + A.OneOf([ + A.Affine(rotate=(-25,25), shear={'x':(-8,8),'y':(-8,8)}, scale=(0.9,1.1), translate_percent={"x":(-0.1,0.1),"y":(-0.1,0.1)}), + A.NoOp()], + p=0.5 + ), + A.HorizontalFlip(p=0.5), + A.VerticalFlip(p=0.5), + A.RandomBrightnessContrast(brightness_limit=0.18, contrast_limit=0.15, p=0.5), + A.HueSaturationValue(hue_shift_limit=7, sat_shift_limit=15, val_shift_limit=10, p=0.5), + A.GaussianNoise(var_limit=(10.0, 30.0), p=0.5), + A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.), + ToTensorV2(transpose_mask=True), + ] + return A.Compose(aug) + else: + aug = [ + A.Resize(IMG_SIZE, IMG_SIZE), + A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.), + ToTensorV2(transpose_mask=True), + ] + return A.Compose(aug) + + print("Section: Dataset and DataLoader Construction") + + class CactusDataset(Dataset): + def __init__(self, img_ids, img_dir, labels=None, transform=None, cache=False): + self.img_ids = img_ids + self.img_dir = img_dir + self.labels = labels # None for test + self.transform = transform + self.cache = cache + self._cache = {} + def __len__(self): + return len(self.img_ids) + def __getitem__(self, idx): + img_id = self.img_ids[idx] + if self.cache and img_id in self._cache: + img4 = self._cache[img_id] + else: + img_path = os.path.join(self.img_dir, img_id) + img4 = load_img_as_numpy_with_mask(img_path) + if self.cache: + self._cache[img_id] = img4 + transformed = self.transform(image=img4) + img = transformed['image'] + if self.labels is not None: + label = float(self.labels[idx]) + return img, label + else: + return img, img_id + + split_seed = 42 + splitter = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=split_seed) + try: + split = next(splitter.split(train_df['id'], train_df['has_cactus'])) + tr_indices, val_indices = split + except Exception as e: + print(f'Stratified split failed ({e}), falling back to random split') + indices = np.arange(len(train_df)) + np.random.shuffle(indices) + n_val = int(0.2 * len(train_df)) + val_indices = indices[:n_val] + tr_indices = indices[n_val:] + + # Sampling, only in debug mode: sample *after* split + if DEBUG: + tr_sample_size = max(2, int(0.1 * len(tr_indices))) + val_sample_size = max(2, int(0.1 * len(val_indices))) + tr_indices = np.random.choice(tr_indices, tr_sample_size, replace=False) + val_indices = np.random.choice(val_indices, val_sample_size, replace=False) + + tr_ids = train_df.iloc[tr_indices]['id'].tolist() + val_ids = train_df.iloc[val_indices]['id'].tolist() + tr_lbls = train_df.iloc[tr_indices]['has_cactus'].tolist() + val_lbls = train_df.iloc[val_indices]['has_cactus'].tolist() + + # For reproducibility and fast debug, cache only in debug for train/val. + train_ds = CactusDataset(tr_ids, TRAIN_DIR, tr_lbls, transform=get_transforms('train'), cache=(DEBUG)) + val_ds = CactusDataset(val_ids, TRAIN_DIR, val_lbls, transform=get_transforms('val'), cache=(DEBUG)) + test_ds = CactusDataset(test_ids, TEST_DIR, labels=None, transform=get_transforms('val'), cache=False) + + BATCH_SIZE = 32 if not DEBUG else 8 + NUM_WORKERS = min(4, os.cpu_count()) + + train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, drop_last=False, num_workers=NUM_WORKERS, pin_memory=True) + val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, drop_last=False, num_workers=NUM_WORKERS, pin_memory=True) + test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE*2, shuffle=False, drop_last=False, num_workers=NUM_WORKERS, pin_memory=True) + + print("Section: Model Definition and Adaptation") + class EfficientNetB0_4ch(nn.Module): + def __init__(self, pretrained=True): + super().__init__() + from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights + if pretrained: + wts = EfficientNet_B0_Weights.DEFAULT + net = efficientnet_b0(weights=wts) + else: + net = efficientnet_b0(weights=None) + old_conv = net.features[0][0] + new_conv = nn.Conv2d(4, old_conv.out_channels, kernel_size=old_conv.kernel_size, + stride=old_conv.stride, padding=old_conv.padding, bias=False) + with torch.no_grad(): + new_conv.weight[:, :3] = old_conv.weight + mean_wt = torch.mean(old_conv.weight, dim=1, keepdim=True) + new_conv.weight[:, 3:4] = mean_wt + net.features[0][0] = new_conv + self.features = net.features + self.avgpool = net.avgpool + inner_dim = net.classifier[1].in_features + self.head = nn.Sequential( + nn.Dropout(0.3), + nn.Linear(inner_dim, 1) + ) + def forward(self, x): + x = self.features(x) + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.head(x) + return x + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + MODEL_TRAINED_FILE = os.path.join(MODEL_DIR, 'efficientnet_b0_best.pth') + scaler = torch.cuda.amp.GradScaler() if torch.cuda.is_available() else None + + # Timing stats for debug regardless path + debug_time = None + estimated_time = None + + NEED_TRAIN = not (os.path.isfile(MODEL_TRAINED_FILE)) + if not NEED_TRAIN: + print("Model checkpoint detected, will use it for inference!") + model = EfficientNetB0_4ch(pretrained=False).to(device) + state = torch.load(MODEL_TRAINED_FILE, map_location=device) + model.load_state_dict(state['model']) + # If in debug, set fake small debug_time for inference-only, as required for compliance. + if DEBUG: + debug_time = 1.0 + scale = (1/0.1) * (1 if DEBUG else 20) + estimated_time = debug_time * scale + else: + print("Model checkpoint not found, proceeding to training...") + print("Section: Training: Staged Fine-Tuning with Discriminative LRs") + model = EfficientNetB0_4ch(pretrained=True).to(device) + criterion = nn.BCEWithLogitsLoss() + backbone_params = [] + mid_params = [] + head_params = list(model.head.parameters()) + for i, m in enumerate(model.features): + if i <= 2: + backbone_params += list(m.parameters()) + elif 3 <= i <= 5: + mid_params += list(m.parameters()) + def set_requires_grad(modules, req): + for m in modules: + for param in m.parameters(): + param.requires_grad = req + set_requires_grad([model.features], False) + set_requires_grad([model.head], True) + EPOCHS = 20 if not DEBUG else 1 + patience = 5 + optimizer = optim.Adam(model.head.parameters(), lr=5e-4, weight_decay=1e-5) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) + best_loss = float('inf') + best_state = None + patience_counter = 0 + start_time = time.time() if DEBUG else None + for epoch in range(EPOCHS): + print(f"Epoch {epoch+1}/{EPOCHS}") + if epoch == 3: + set_requires_grad([model.features[3], model.features[4], model.features[5]], True) + optimizer = optim.Adam([ + {'params': backbone_params, 'lr': 1e-4}, + {'params': mid_params, 'lr': 2e-4}, + {'params': head_params, 'lr':5e-4}, + ], weight_decay=1e-5) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS-epoch) + print("Unfroze mid layers of EfficientNet for fine-tuning.") + elif epoch == 6: + set_requires_grad([model.features], True) + print("Unfroze all layers of EfficientNet for full fine-tuning.") + + model.train() + tr_loss = 0. + tr_cnt = 0 + for imgs, lbls in train_loader: + imgs = imgs.to(device) + lbls = lbls.to(device).view(-1,1) + optimizer.zero_grad() + if scaler is not None: + with torch.cuda.amp.autocast(): + outs = model(imgs) + loss = criterion(outs, lbls) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + outs = model(imgs) + loss = criterion(outs, lbls) + loss.backward() + optimizer.step() + tr_loss += loss.item() * imgs.size(0) + tr_cnt += imgs.size(0) + if scheduler is not None: + scheduler.step() + + tr_loss = tr_loss / tr_cnt + + model.eval() + val_loss = 0. + val_cnt = 0 + all_val_lbls = [] + all_val_preds = [] + with torch.no_grad(): + for imgs, lbls in val_loader: + imgs = imgs.to(device) + lbls = lbls.cpu().numpy() + outs = model(imgs).cpu().squeeze().numpy() + preds = 1/(1 + np.exp(-outs)) + loss = criterion(torch.tensor(outs).view(-1,1), torch.tensor(lbls).view(-1,1)).item() + val_loss += loss * imgs.size(0) + val_cnt += imgs.size(0) + all_val_lbls.append(lbls) + all_val_preds.append(preds) + val_loss = val_loss / val_cnt + all_val_lbls = np.concatenate(all_val_lbls) + all_val_preds = np.concatenate(all_val_preds) + try: + val_logloss = log_loss(all_val_lbls, all_val_preds, eps=1e-7) + except Exception as ex: + val_logloss = float('inf') + print("Error computing log_loss on val:", ex) + + print(f"Train Loss: {tr_loss:.5f} | Val Loss (BCE): {val_loss:.5f} | Val LogLoss: {val_logloss:.5f}") + + if val_logloss < best_loss: + best_loss = val_logloss + best_state = { + 'model': model.state_dict(), + 'epoch': epoch, + 'val_loss': best_loss, + } + torch.save(best_state, MODEL_TRAINED_FILE) + patience_counter = 0 + print(f"Best model saved. (epoch {epoch+1}, val_logloss={val_logloss:.5f})") + else: + patience_counter += 1 + print(f"No improvement. Early stopping patience: {patience_counter}/{patience}") + + if patience_counter >= patience: + print(f"Early stopping triggered at epoch {epoch+1}.") + break + if DEBUG and start_time is not None: + end_time = time.time() + debug_time = end_time - start_time + # Compute estimated time: (fractional data)*(epochs) compared + sample_factor = 0.1 + scale = (1/sample_factor) * (20 if not DEBUG else 1) + estimated_time = debug_time * scale + # Reload best model for evaluation + state = torch.load(MODEL_TRAINED_FILE, map_location=device) + model.load_state_dict(state['model']) + + print("Section: Validation Evaluation and Metric Calculation") + model.eval() + val_lbls, val_prs = [], [] + with torch.no_grad(): + for imgs, lbls in val_loader: + imgs = imgs.to(device) + outs = model(imgs).cpu().squeeze().numpy() + prs = 1/(1+np.exp(-outs)) + val_lbls.append(lbls.numpy()) + val_prs.append(prs) + val_lbls = np.concatenate(val_lbls) + val_prs = np.concatenate(val_prs) + try: + val_logloss = log_loss(val_lbls, val_prs, eps=1e-7) + except Exception as ex: + val_logloss = float('inf') + print("Error computing log_loss on validation:", ex) + print(f"Final best model log loss on validation split: {val_logloss:.6f}") + scores = pd.DataFrame( + {'Model': ['efficientnet_b0', 'ensemble'], 'LogLoss': [val_logloss, val_logloss]} + ).set_index('Model') + scores.to_csv(SCORES_PATH) + print(f"Saved scores.csv with validation log loss.") + + print("Section: Prediction and Submission Generation") + model.eval() + test_probs = [] + test_ids_ordered = [] + with torch.no_grad(): + for imgs, img_ids in test_loader: + imgs = imgs.to(device) + outs = model(imgs).cpu().squeeze().numpy() + prs = 1/(1+np.exp(-outs)) + if isinstance(img_ids, list) or isinstance(img_ids, np.ndarray): + test_ids_ordered += list(img_ids) + else: + test_ids_ordered.append(img_ids) + test_probs.extend(np.array(prs).ravel().tolist()) + submit_df = pd.DataFrame({'id': test_ids_ordered, 'has_cactus': test_probs}) + submit_df = submit_df.set_index('id') + try: + submit_df = submit_df.reindex(sample_submission['id']).reset_index() + except Exception: + submit_df = submit_df.reset_index() + submit_df['has_cactus'] = submit_df['has_cactus'].clip(0,1) + submit_df.to_csv(SUBMISSION_PATH, index=False, float_format='%.6f') + print(f"Saved submission.csv with {len(submit_df)} rows. Format: {submit_df.columns.tolist()}") + + # === Debug info output, always print in debug mode, even if only inference === + if DEBUG: + if debug_time is None: + debug_time = 1.0 + scale = (1/0.1)*(1 if DEBUG else 20) + estimated_time = debug_time * scale + print("=== Start of Debug Information ===") + print(f"debug_time: {debug_time}") + print(f"estimated_time: {estimated_time}") + print("=== End of Debug Information ===") + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/test/notebook/testfiles/main_missing_main_fn.py b/test/notebook/testfiles/main_missing_main_fn.py new file mode 100644 index 00000000..d4289ed5 --- /dev/null +++ b/test/notebook/testfiles/main_missing_main_fn.py @@ -0,0 +1,507 @@ +import os +import sys +import time +import random +import numpy as np +import pandas as pd + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader + +import timm +import albumentations as A +from albumentations.pytorch import ToTensorV2 + +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import roc_auc_score, confusion_matrix + +import cv2 +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument('--debug', action='store_true', help='Run in debug mode') +args = parser.parse_args() +DEBUG = args.debug + +SEED = 2024 +np.random.seed(SEED) +random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +TRAIN_DIR = './workspace_input/train/' +TEST_DIR = './workspace_input/test/' +TRAIN_CSV = './workspace_input/train.csv' +SAMPLE_SUB_PATH = './workspace_input/sample_submission.csv' +MODEL_DIR = 'models/' +os.makedirs(MODEL_DIR, exist_ok=True) + +def print_eda(train_df): + print("=== Start of EDA part ===") + print("Shape of train.csv:", train_df.shape) + print("First 5 rows:\n", train_df.head()) + print("Column data types:\n", train_df.dtypes) + print("Missing values per column:\n", train_df.isnull().sum()) + print("Unique values per column:") + for col in train_df.columns: + print(f" - {col}: {train_df[col].nunique()}") + label_counts = train_df['has_cactus'].value_counts() + print("Label distribution (has_cactus):") + print(label_counts) + pos, neg = label_counts.get(1, 0), label_counts.get(0, 0) + total = pos + neg + if total > 0: + print(f" Positive:Negative ratio: {pos}:{neg} ({pos/total:.3f}:{neg/total:.3f})") + print(f" Percentage positive: {pos/total*100:.2f}%") + else: + print(" No data found.") + print("Image filename examples:", train_df['id'].unique()[:5]) + print("=== End of EDA part ===") + +class CactusDataset(Dataset): + def __init__(self, image_ids, labels=None, id2path=None, transforms=None): + self.image_ids = image_ids + self.labels = labels + self.id2path = id2path + self.transforms = transforms + + def __len__(self): + return len(self.image_ids) + + def __getitem__(self, idx): + img_id = self.image_ids[idx] + img_path = self.id2path[img_id] + image = cv2.imread(img_path) + if image is None: + raise RuntimeError(f"Cannot read image at {img_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + if self.transforms: + augmented = self.transforms(image=image) + image = augmented["image"] + if self.labels is not None: + label = self.labels[idx] + return image, label, img_id + else: + return image, img_id + +def get_transforms(mode='train'): + # Correct Cutout: Albumentations v1.4.15 provides 'Cutout' as a class, but not always in the root. + # Defensive import; fallback to the most robust method for v1.4.15 + imagenet_mean = [0.485, 0.456, 0.406] + imagenet_std = [0.229, 0.224, 0.225] + if mode == 'train': + min_frac, max_frac = 0.05, 0.2 + min_cut = int(300 * min_frac) + max_cut = int(300 * max_frac) + # There is no A.Cutout in v1.4.15 root, but A.augmentations.transforms.Cutout exists. + try: + from albumentations.augmentations.transforms import Cutout + have_cutout = True + except ImportError: + have_cutout = False + this_cut_h = random.randint(min_cut, max_cut) + this_cut_w = random.randint(min_cut, max_cut) + cutout_fill = [int(255 * m) for m in imagenet_mean] + tforms = [ + A.RandomResizedCrop(300, 300, scale=(0.7, 1.0), ratio=(0.8, 1.2), p=1.0), + A.Rotate(limit=30, p=0.8), + ] + if have_cutout: + tforms.append( + Cutout( + num_holes=1, + max_h_size=this_cut_h, + max_w_size=this_cut_w, + fill_value=cutout_fill, # RGB image in albumentations requires [R,G,B] + always_apply=False, + p=0.7 + ) + ) + else: + # No available Cutout, so fallback to no cutout but emit warning + print("WARNING: albumentations.Cutout not found, continuing without Cutout augmentation") + tforms.extend([ + A.RandomContrast(limit=0.2, p=0.5), + A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.1), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + return A.Compose(tforms) + else: + return A.Compose([ + A.Resize(300, 300), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + +def get_dataloader(dataset, batch_size, shuffle=False, num_workers=4, pin_memory=True): + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=pin_memory + ) + +def get_efficientnet_b3(dropout_rate=0.3): + model = timm.create_model('efficientnet_b3', pretrained=True) + n_in = model.classifier.in_features if hasattr(model, "classifier") else model.fc.in_features + model.classifier = nn.Sequential( + nn.Dropout(dropout_rate), + nn.Linear(n_in, 1) + ) + return model + +def compute_class_weight(y): + counts = np.bincount(y) + if len(counts) < 2: + counts = np.pad(counts, (0, 2-len(counts)), constant_values=0) + n_pos, n_neg = counts[1], counts[0] + total = n_pos + n_neg + minority, majority = min(n_pos, n_neg), max(n_pos, n_neg) + ratio = majority / (minority + 1e-10) + need_weights = ratio > 2 + weights = None + if need_weights: + inv_freq = [1 / (n_neg + 1e-10), 1 / (n_pos + 1e-10)] + s = sum(inv_freq) + weights = [w / s * 2 for w in inv_freq] + return weights, n_pos, n_neg, ratio, need_weights + +def train_one_epoch(model, loss_fn, optimizer, scheduler, dataloader, device, class_weights): + model.train() + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + if scheduler is not None: + scheduler.step() + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + avg_loss = total_loss / total_samples + return avg_loss + +@torch.no_grad() +def eval_model(model, loss_fn, dataloader, device, class_weights): + model.eval() + y_true, y_pred = [], [] + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + probs = torch.sigmoid(logits) + y_true.append(labels.cpu().numpy()) + y_pred.append(probs.cpu().numpy()) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + y_true = np.vstack(y_true).reshape(-1) + y_pred = np.vstack(y_pred).reshape(-1) + avg_loss = total_loss / total_samples + return avg_loss, y_true, y_pred + +def confusion_info(y_true, y_pred, threshold=0.5): + preds = (y_pred > threshold).astype(int) + cm = confusion_matrix(y_true, preds) + return cm + +def inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, class_weights, need_weights, + BATCH_SIZE, N_WORKERS, cv_fold): + oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] + for fold in range(cv_fold): + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + val_ds = CactusDataset(val_img_ids, val_labels, id2path=train_id2path, transforms=get_transforms("val")) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(val_auc) + print(f"Reloaded fold {fold}, OOF Validation AUC={val_auc:.5f}") + + all_oof_true = np.concatenate(oof_true) + all_oof_pred = np.concatenate(oof_pred) + oof_auc = roc_auc_score(all_oof_true, all_oof_pred) + oof_cm = confusion_info(all_oof_true, all_oof_pred) + print(f"OOF ROC-AUC (from loaded models): {oof_auc:.5f}") + print(f"OOF Confusion Matrix:\n{oof_cm}") + + test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") + ) + test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + test_pred_list = [] + for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") + test_probs = np.mean(test_pred_list, axis=0) + + submission = pd.read_csv(SAMPLE_SUB_PATH) + submission['has_cactus'] = test_probs + submission.to_csv('submission.csv', index=False) + print(f"Saved submission.csv in required format with {len(submission)} rows.") + + scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] + }) + scores_df.set_index('Model', inplace=True) + scores_df.to_csv("scores.csv") + print(f"Saved cross-validation scores to scores.csv") + +print("Section: Data Loading and Preprocessing") +try: + train_df = pd.read_csv(TRAIN_CSV) +except Exception as e: + print(f"Failed to load train.csv: {e}") + sys.exit(1) +print_eda(train_df) + +train_id2path = {img_id: os.path.join(TRAIN_DIR, img_id) for img_id in train_df['id']} +try: + sample_sub = pd.read_csv(SAMPLE_SUB_PATH) +except Exception as e: + print(f"Failed to load sample_submission.csv: {e}") + sys.exit(1) +test_img_ids = list(sample_sub['id']) +test_id2path = {img_id: os.path.join(TEST_DIR, img_id) for img_id in test_img_ids} +print(f"Loaded {len(train_id2path)} train images, {len(test_id2path)} test images.") + +y_train = train_df['has_cactus'].values +class_weights, n_pos, n_neg, imbalance_ratio, need_weights = compute_class_weight(y_train) +print(f"Class stats: Pos={n_pos}, Neg={n_neg}, Imbalance Ratio(majority/minority)={imbalance_ratio:.3f}") +print(f"Use class weights: {need_weights}, Class weights: {class_weights if class_weights is not None else '[1.0,1.0]'}") +if class_weights is not None: + np.save(os.path.join(MODEL_DIR, "class_weights.npy"), class_weights) + +print("Section: Feature Engineering") +train_df = train_df.copy() +cv_fold = 5 +skf = StratifiedKFold(n_splits=cv_fold, shuffle=True, random_state=SEED) +folds = np.zeros(len(train_df), dtype=np.int32) +for idx, (_, val_idx) in enumerate(skf.split(train_df['id'], train_df['has_cactus'])): + folds[val_idx] = idx +train_df['fold'] = folds +print(f"Assigned stratified {cv_fold}-fold indices. Fold sample counts:") +for f in range(cv_fold): + dist = train_df.loc[train_df['fold'] == f, 'has_cactus'].value_counts().to_dict() + print(f" Fold {f}: n={len(train_df[train_df['fold'] == f])} class dist={dist}") + +print("Section: Model Training and Evaluation") +dropout_rate = round(random.uniform(0.2, 0.5), 2) +print(f"Model config: EfficientNet-B3, Image size 300, Head dropout={dropout_rate}") + +if DEBUG: + print("DEBUG mode: using 10% subsample and 1 epoch (per fold)") + sample_frac = 0.10 + sampled_idxs = [] + for f in range(cv_fold): + fold_idx = train_df.index[train_df['fold'] == f].tolist() + fold_labels = train_df.loc[fold_idx, 'has_cactus'].values + idx_pos = [i for i, l in zip(fold_idx, fold_labels) if l == 1] + idx_neg = [i for i, l in zip(fold_idx, fold_labels) if l == 0] + n_pos = max(1, int(sample_frac * len(idx_pos))) + n_neg = max(1, int(sample_frac * len(idx_neg))) + if len(idx_pos) > 0: + sampled_idxs += np.random.choice(idx_pos, n_pos, replace=False).tolist() + if len(idx_neg) > 0: + sampled_idxs += np.random.choice(idx_neg, n_neg, replace=False).tolist() + train_df = train_df.loc[sampled_idxs].reset_index(drop=True) + print(f"DEBUG subsample shape: {train_df.shape}") + debug_epochs = 1 +else: + debug_epochs = None + +BATCH_SIZE = 64 if torch.cuda.is_available() else 32 +N_WORKERS = 4 if torch.cuda.is_available() else 1 +EPOCHS = 20 if not DEBUG else debug_epochs +MIN_EPOCHS = 5 if not DEBUG else 1 +EARLY_STOP_PATIENCE = 7 if not DEBUG else 2 +LR = 1e-3 + +model_files = [os.path.join(MODEL_DIR, f"efficientnet_b3_fold{f}.pt") for f in range(cv_fold)] +if all([os.path.exists(f) for f in model_files]): + print("All fold models found in models/. Running inference and file saving only (no retrain).") + inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, + class_weights, need_weights, BATCH_SIZE, N_WORKERS, cv_fold) + return + +oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] +start_time = time.time() if DEBUG else None + +for fold in range(cv_fold): + print(f"\n=== FOLD {fold} TRAINING ===") + df_train = train_df[train_df['fold'] != fold].reset_index(drop=True) + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + print(f"Train size: {df_train.shape[0]}, Val size: {df_val.shape[0]}") + train_img_ids = df_train['id'].tolist() + train_labels = df_train['has_cactus'].values + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + + train_ds = CactusDataset( + train_img_ids, train_labels, + id2path=train_id2path, + transforms=get_transforms("train") + ) + val_ds = CactusDataset( + val_img_ids, val_labels, + id2path=train_id2path, + transforms=get_transforms("val") + ) + train_loader = get_dataloader(train_ds, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + optimizer = optim.AdamW(model.parameters(), lr=LR) + scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + best_auc = -np.inf + best_epoch = -1 + best_model_state = None + patience = 0 + + for epoch in range(EPOCHS): + train_loss = train_one_epoch( + model, loss_fn, optimizer, scheduler, train_loader, DEVICE, fold_class_weights) + val_loss, val_true, val_pred = eval_model( + model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + cm = confusion_info(val_true, val_pred) + print(f"Epoch {epoch+1:02d}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_auc={val_auc:.4f}") + print(f" Val confusion_matrix (rows:true [0,1]; cols:pred [0,1]):\n{cm}") + if val_auc > best_auc: + best_auc = val_auc + best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} + best_epoch = epoch + patience = 0 + else: + patience += 1 + if DEBUG and epoch + 1 >= debug_epochs: + break + if (epoch + 1) >= MIN_EPOCHS and patience >= EARLY_STOP_PATIENCE: + print(f"Early stopping at epoch {epoch+1}, best_epoch={best_epoch+1}.") + break + + model.load_state_dict(best_model_state) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + torch.save(model.state_dict(), fold_model_path) + print(f"Saved best model for fold {fold} at {fold_model_path} (best_auc={best_auc:.5f}, best_epoch={best_epoch+1})") + + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(best_auc) + print(f"OOF stored for fold {fold}, Validation AUC={best_auc:.5f}") + +end_time = time.time() if DEBUG else None +if DEBUG: + debug_time = end_time - start_time + estimated_time = (1 / 0.1) * (EPOCHS / debug_epochs) * debug_time + print("=== Start of Debug Information ===") + print(f"debug_time: {debug_time:.1f}") + print(f"estimated_time: {estimated_time:.1f}") + print("=== End of Debug Information ===") + +print("\nSection: Ensemble Strategy and Final Predictions") +all_oof_true = np.concatenate(oof_true) +all_oof_pred = np.concatenate(oof_pred) +oof_auc = roc_auc_score(all_oof_true, all_oof_pred) +oof_cm = confusion_info(all_oof_true, all_oof_pred) +print(f"OOF ROC-AUC: {oof_auc:.5f}") +print(f"OOF Confusion Matrix:\n{oof_cm}") + +test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") +) +test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) +test_pred_list = [] +for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") +test_probs = np.mean(test_pred_list, axis=0) + +print("Section: Submission File Generation") +submission = pd.read_csv(SAMPLE_SUB_PATH) +submission['has_cactus'] = test_probs +submission.to_csv('submission.csv', index=False) +print(f"Saved submission.csv in required format with {len(submission)} rows.") + +scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] +}) +scores_df.set_index('Model', inplace=True) +scores_df.to_csv("scores.csv") +print(f"Saved cross-validation scores to scores.csv") diff --git a/test/notebook/testfiles/main_missing_sections.py b/test/notebook/testfiles/main_missing_sections.py new file mode 100644 index 00000000..009a1fcf --- /dev/null +++ b/test/notebook/testfiles/main_missing_sections.py @@ -0,0 +1,506 @@ +import os +import sys +import time +import random +import numpy as np +import pandas as pd + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader + +import timm +import albumentations as A +from albumentations.pytorch import ToTensorV2 + +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import roc_auc_score, confusion_matrix + +import cv2 +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument('--debug', action='store_true', help='Run in debug mode') +args = parser.parse_args() +DEBUG = args.debug + +SEED = 2024 +np.random.seed(SEED) +random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +TRAIN_DIR = './workspace_input/train/' +TEST_DIR = './workspace_input/test/' +TRAIN_CSV = './workspace_input/train.csv' +SAMPLE_SUB_PATH = './workspace_input/sample_submission.csv' +MODEL_DIR = 'models/' +os.makedirs(MODEL_DIR, exist_ok=True) + +def print_eda(train_df): + print("=== Start of EDA part ===") + print("Shape of train.csv:", train_df.shape) + print("First 5 rows:\n", train_df.head()) + print("Column data types:\n", train_df.dtypes) + print("Missing values per column:\n", train_df.isnull().sum()) + print("Unique values per column:") + for col in train_df.columns: + print(f" - {col}: {train_df[col].nunique()}") + label_counts = train_df['has_cactus'].value_counts() + print("Label distribution (has_cactus):") + print(label_counts) + pos, neg = label_counts.get(1, 0), label_counts.get(0, 0) + total = pos + neg + if total > 0: + print(f" Positive:Negative ratio: {pos}:{neg} ({pos/total:.3f}:{neg/total:.3f})") + print(f" Percentage positive: {pos/total*100:.2f}%") + else: + print(" No data found.") + print("Image filename examples:", train_df['id'].unique()[:5]) + print("=== End of EDA part ===") + +class CactusDataset(Dataset): + def __init__(self, image_ids, labels=None, id2path=None, transforms=None): + self.image_ids = image_ids + self.labels = labels + self.id2path = id2path + self.transforms = transforms + + def __len__(self): + return len(self.image_ids) + + def __getitem__(self, idx): + img_id = self.image_ids[idx] + img_path = self.id2path[img_id] + image = cv2.imread(img_path) + if image is None: + raise RuntimeError(f"Cannot read image at {img_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + if self.transforms: + augmented = self.transforms(image=image) + image = augmented["image"] + if self.labels is not None: + label = self.labels[idx] + return image, label, img_id + else: + return image, img_id + +def get_transforms(mode='train'): + # Correct Cutout: Albumentations v1.4.15 provides 'Cutout' as a class, but not always in the root. + # Defensive import; fallback to the most robust method for v1.4.15 + imagenet_mean = [0.485, 0.456, 0.406] + imagenet_std = [0.229, 0.224, 0.225] + if mode == 'train': + min_frac, max_frac = 0.05, 0.2 + min_cut = int(300 * min_frac) + max_cut = int(300 * max_frac) + # There is no A.Cutout in v1.4.15 root, but A.augmentations.transforms.Cutout exists. + try: + from albumentations.augmentations.transforms import Cutout + have_cutout = True + except ImportError: + have_cutout = False + this_cut_h = random.randint(min_cut, max_cut) + this_cut_w = random.randint(min_cut, max_cut) + cutout_fill = [int(255 * m) for m in imagenet_mean] + tforms = [ + A.RandomResizedCrop(300, 300, scale=(0.7, 1.0), ratio=(0.8, 1.2), p=1.0), + A.Rotate(limit=30, p=0.8), + ] + if have_cutout: + tforms.append( + Cutout( + num_holes=1, + max_h_size=this_cut_h, + max_w_size=this_cut_w, + fill_value=cutout_fill, # RGB image in albumentations requires [R,G,B] + always_apply=False, + p=0.7 + ) + ) + else: + # No available Cutout, so fallback to no cutout but emit warning + print("WARNING: albumentations.Cutout not found, continuing without Cutout augmentation") + tforms.extend([ + A.RandomContrast(limit=0.2, p=0.5), + A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.1), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + return A.Compose(tforms) + else: + return A.Compose([ + A.Resize(300, 300), + A.Normalize(mean=imagenet_mean, std=imagenet_std, max_pixel_value=255.0), + ToTensorV2() + ]) + +def get_dataloader(dataset, batch_size, shuffle=False, num_workers=4, pin_memory=True): + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=pin_memory + ) + +def get_efficientnet_b3(dropout_rate=0.3): + model = timm.create_model('efficientnet_b3', pretrained=True) + n_in = model.classifier.in_features if hasattr(model, "classifier") else model.fc.in_features + model.classifier = nn.Sequential( + nn.Dropout(dropout_rate), + nn.Linear(n_in, 1) + ) + return model + +def compute_class_weight(y): + counts = np.bincount(y) + if len(counts) < 2: + counts = np.pad(counts, (0, 2-len(counts)), constant_values=0) + n_pos, n_neg = counts[1], counts[0] + total = n_pos + n_neg + minority, majority = min(n_pos, n_neg), max(n_pos, n_neg) + ratio = majority / (minority + 1e-10) + need_weights = ratio > 2 + weights = None + if need_weights: + inv_freq = [1 / (n_neg + 1e-10), 1 / (n_pos + 1e-10)] + s = sum(inv_freq) + weights = [w / s * 2 for w in inv_freq] + return weights, n_pos, n_neg, ratio, need_weights + +def train_one_epoch(model, loss_fn, optimizer, scheduler, dataloader, device, class_weights): + model.train() + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + if scheduler is not None: + scheduler.step() + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + avg_loss = total_loss / total_samples + return avg_loss + +@torch.no_grad() +def eval_model(model, loss_fn, dataloader, device, class_weights): + model.eval() + y_true, y_pred = [], [] + total_loss = 0.0 + total_samples = 0 + for batch in dataloader: + images, labels, _ = batch + images = images.to(device) + labels = labels.float().unsqueeze(1).to(device) + logits = model(images) + probs = torch.sigmoid(logits) + y_true.append(labels.cpu().numpy()) + y_pred.append(probs.cpu().numpy()) + if class_weights is not None: + weight = labels * class_weights[1] + (1 - labels) * class_weights[0] + loss = loss_fn(logits, labels) + loss = (loss * weight).mean() + else: + loss = loss_fn(logits, labels) + total_loss += loss.item() * labels.size(0) + total_samples += labels.size(0) + y_true = np.vstack(y_true).reshape(-1) + y_pred = np.vstack(y_pred).reshape(-1) + avg_loss = total_loss / total_samples + return avg_loss, y_true, y_pred + +def confusion_info(y_true, y_pred, threshold=0.5): + preds = (y_pred > threshold).astype(int) + cm = confusion_matrix(y_true, preds) + return cm + +def inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, class_weights, need_weights, + BATCH_SIZE, N_WORKERS, cv_fold): + oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] + for fold in range(cv_fold): + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + val_ds = CactusDataset(val_img_ids, val_labels, id2path=train_id2path, transforms=get_transforms("val")) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(val_auc) + print(f"Reloaded fold {fold}, OOF Validation AUC={val_auc:.5f}") + + all_oof_true = np.concatenate(oof_true) + all_oof_pred = np.concatenate(oof_pred) + oof_auc = roc_auc_score(all_oof_true, all_oof_pred) + oof_cm = confusion_info(all_oof_true, all_oof_pred) + print(f"OOF ROC-AUC (from loaded models): {oof_auc:.5f}") + print(f"OOF Confusion Matrix:\n{oof_cm}") + + test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") + ) + test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + test_pred_list = [] + for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") + test_probs = np.mean(test_pred_list, axis=0) + + submission = pd.read_csv(SAMPLE_SUB_PATH) + submission['has_cactus'] = test_probs + submission.to_csv('submission.csv', index=False) + print(f"Saved submission.csv in required format with {len(submission)} rows.") + + scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] + }) + scores_df.set_index('Model', inplace=True) + scores_df.to_csv("scores.csv") + print(f"Saved cross-validation scores to scores.csv") + +def main(): + try: + train_df = pd.read_csv(TRAIN_CSV) + except Exception as e: + print(f"Failed to load train.csv: {e}") + sys.exit(1) + print_eda(train_df) + + train_id2path = {img_id: os.path.join(TRAIN_DIR, img_id) for img_id in train_df['id']} + try: + sample_sub = pd.read_csv(SAMPLE_SUB_PATH) + except Exception as e: + print(f"Failed to load sample_submission.csv: {e}") + sys.exit(1) + test_img_ids = list(sample_sub['id']) + test_id2path = {img_id: os.path.join(TEST_DIR, img_id) for img_id in test_img_ids} + print(f"Loaded {len(train_id2path)} train images, {len(test_id2path)} test images.") + + y_train = train_df['has_cactus'].values + class_weights, n_pos, n_neg, imbalance_ratio, need_weights = compute_class_weight(y_train) + print(f"Class stats: Pos={n_pos}, Neg={n_neg}, Imbalance Ratio(majority/minority)={imbalance_ratio:.3f}") + print(f"Use class weights: {need_weights}, Class weights: {class_weights if class_weights is not None else '[1.0,1.0]'}") + if class_weights is not None: + np.save(os.path.join(MODEL_DIR, "class_weights.npy"), class_weights) + + train_df = train_df.copy() + cv_fold = 5 + skf = StratifiedKFold(n_splits=cv_fold, shuffle=True, random_state=SEED) + folds = np.zeros(len(train_df), dtype=np.int32) + for idx, (_, val_idx) in enumerate(skf.split(train_df['id'], train_df['has_cactus'])): + folds[val_idx] = idx + train_df['fold'] = folds + print(f"Assigned stratified {cv_fold}-fold indices. Fold sample counts:") + for f in range(cv_fold): + dist = train_df.loc[train_df['fold'] == f, 'has_cactus'].value_counts().to_dict() + print(f" Fold {f}: n={len(train_df[train_df['fold'] == f])} class dist={dist}") + + dropout_rate = round(random.uniform(0.2, 0.5), 2) + print(f"Model config: EfficientNet-B3, Image size 300, Head dropout={dropout_rate}") + + if DEBUG: + print("DEBUG mode: using 10% subsample and 1 epoch (per fold)") + sample_frac = 0.10 + sampled_idxs = [] + for f in range(cv_fold): + fold_idx = train_df.index[train_df['fold'] == f].tolist() + fold_labels = train_df.loc[fold_idx, 'has_cactus'].values + idx_pos = [i for i, l in zip(fold_idx, fold_labels) if l == 1] + idx_neg = [i for i, l in zip(fold_idx, fold_labels) if l == 0] + n_pos = max(1, int(sample_frac * len(idx_pos))) + n_neg = max(1, int(sample_frac * len(idx_neg))) + if len(idx_pos) > 0: + sampled_idxs += np.random.choice(idx_pos, n_pos, replace=False).tolist() + if len(idx_neg) > 0: + sampled_idxs += np.random.choice(idx_neg, n_neg, replace=False).tolist() + train_df = train_df.loc[sampled_idxs].reset_index(drop=True) + print(f"DEBUG subsample shape: {train_df.shape}") + debug_epochs = 1 + else: + debug_epochs = None + + BATCH_SIZE = 64 if torch.cuda.is_available() else 32 + N_WORKERS = 4 if torch.cuda.is_available() else 1 + EPOCHS = 20 if not DEBUG else debug_epochs + MIN_EPOCHS = 5 if not DEBUG else 1 + EARLY_STOP_PATIENCE = 7 if not DEBUG else 2 + LR = 1e-3 + + model_files = [os.path.join(MODEL_DIR, f"efficientnet_b3_fold{f}.pt") for f in range(cv_fold)] + if all([os.path.exists(f) for f in model_files]): + print("All fold models found in models/. Running inference and file saving only (no retrain).") + inference_and_submission(train_df, train_id2path, test_img_ids, test_id2path, dropout_rate, + class_weights, need_weights, BATCH_SIZE, N_WORKERS, cv_fold) + return + + oof_true, oof_pred, fold_scores, fold_val_ids = [], [], [], [] + start_time = time.time() if DEBUG else None + + for fold in range(cv_fold): + print(f"\n=== FOLD {fold} TRAINING ===") + df_train = train_df[train_df['fold'] != fold].reset_index(drop=True) + df_val = train_df[train_df['fold'] == fold].reset_index(drop=True) + print(f"Train size: {df_train.shape[0]}, Val size: {df_val.shape[0]}") + train_img_ids = df_train['id'].tolist() + train_labels = df_train['has_cactus'].values + val_img_ids = df_val['id'].tolist() + val_labels = df_val['has_cactus'].values + + train_ds = CactusDataset( + train_img_ids, train_labels, + id2path=train_id2path, + transforms=get_transforms("train") + ) + val_ds = CactusDataset( + val_img_ids, val_labels, + id2path=train_id2path, + transforms=get_transforms("val") + ) + train_loader = get_dataloader(train_ds, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS) + val_loader = get_dataloader(val_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.to(DEVICE) + loss_fn = nn.BCEWithLogitsLoss(reduction='none') + optimizer = optim.AdamW(model.parameters(), lr=LR) + scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) + fold_class_weights = class_weights if need_weights else None + if fold_class_weights is not None: + fold_class_weights = torch.tensor(fold_class_weights).float().to(DEVICE) + best_auc = -np.inf + best_epoch = -1 + best_model_state = None + patience = 0 + + for epoch in range(EPOCHS): + train_loss = train_one_epoch( + model, loss_fn, optimizer, scheduler, train_loader, DEVICE, fold_class_weights) + val_loss, val_true, val_pred = eval_model( + model, loss_fn, val_loader, DEVICE, fold_class_weights) + val_auc = roc_auc_score(val_true, val_pred) + cm = confusion_info(val_true, val_pred) + print(f"Epoch {epoch+1:02d}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_auc={val_auc:.4f}") + print(f" Val confusion_matrix (rows:true [0,1]; cols:pred [0,1]):\n{cm}") + if val_auc > best_auc: + best_auc = val_auc + best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} + best_epoch = epoch + patience = 0 + else: + patience += 1 + if DEBUG and epoch + 1 >= debug_epochs: + break + if (epoch + 1) >= MIN_EPOCHS and patience >= EARLY_STOP_PATIENCE: + print(f"Early stopping at epoch {epoch+1}, best_epoch={best_epoch+1}.") + break + + model.load_state_dict(best_model_state) + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + torch.save(model.state_dict(), fold_model_path) + print(f"Saved best model for fold {fold} at {fold_model_path} (best_auc={best_auc:.5f}, best_epoch={best_epoch+1})") + + _, val_true, val_pred = eval_model(model, loss_fn, val_loader, DEVICE, fold_class_weights) + oof_true.append(val_true) + oof_pred.append(val_pred) + fold_val_ids.append(val_img_ids) + fold_scores.append(best_auc) + print(f"OOF stored for fold {fold}, Validation AUC={best_auc:.5f}") + + end_time = time.time() if DEBUG else None + if DEBUG: + debug_time = end_time - start_time + estimated_time = (1 / 0.1) * (EPOCHS / debug_epochs) * debug_time + print("=== Start of Debug Information ===") + print(f"debug_time: {debug_time:.1f}") + print(f"estimated_time: {estimated_time:.1f}") + print("=== End of Debug Information ===") + + all_oof_true = np.concatenate(oof_true) + all_oof_pred = np.concatenate(oof_pred) + oof_auc = roc_auc_score(all_oof_true, all_oof_pred) + oof_cm = confusion_info(all_oof_true, all_oof_pred) + print(f"OOF ROC-AUC: {oof_auc:.5f}") + print(f"OOF Confusion Matrix:\n{oof_cm}") + + test_ds = CactusDataset( + test_img_ids, labels=None, + id2path=test_id2path, + transforms=get_transforms("val") + ) + test_loader = get_dataloader(test_ds, BATCH_SIZE, shuffle=False, num_workers=N_WORKERS) + test_pred_list = [] + for fold in range(cv_fold): + fold_model_path = os.path.join(MODEL_DIR, f"efficientnet_b3_fold{fold}.pt") + model = get_efficientnet_b3(dropout_rate=dropout_rate) + model.load_state_dict(torch.load(fold_model_path, map_location='cpu')) + model.to(DEVICE) + model.eval() + preds = [] + with torch.no_grad(): + for batch in test_loader: + images, img_ids = batch + images = images.to(DEVICE) + logits = model(images) + probs = torch.sigmoid(logits).cpu().numpy().reshape(-1) + preds.append(probs) + fold_test_pred = np.concatenate(preds) + test_pred_list.append(fold_test_pred) + print(f"Loaded fold {fold} for test prediction.") + test_probs = np.mean(test_pred_list, axis=0) + + submission = pd.read_csv(SAMPLE_SUB_PATH) + submission['has_cactus'] = test_probs + submission.to_csv('submission.csv', index=False) + print(f"Saved submission.csv in required format with {len(submission)} rows.") + + scores_df = pd.DataFrame({ + 'Model': [f"efficientnet_b3_fold{f}" for f in range(cv_fold)] + ['ensemble'], + 'ROC-AUC': list(fold_scores) + [oof_auc] + }) + scores_df.set_index('Model', inplace=True) + scores_df.to_csv("scores.csv") + print(f"Saved cross-validation scores to scores.csv") + +if __name__ == "__main__": + main()