feat: script to collect a few factorio yolo dataset

eslint.config.ts

@@ -41,7 +41,7 @@ export default antfu(
     files: [
       'packages/autorio/**/*.ts',
       'packages/tstl-plugin-reload-factorio-mod/example/*.ts',
-      'packages/factorio-rcon-snippets/**/*.ts',
+      'packages/factorio-rcon-snippets-for-vscode/**/*.ts',
     ],
   },
   // #endregion

packages/factorio-rcon-snippets-for-node/README.md

@@ -0,0 +1,9 @@
+# factorio-rcon-snippets-for-node
+
+This package contains some snippets for do some cool tasks.
+
+## Usage
+
+```bash
+RCON_PORT=2302 RCON_HOST=127.0.0.1 RCON_PASSWORD=123456 node src/factorio_yolo_dataset_collector_v0.ts
+```

packages/factorio-rcon-snippets-for-node/package.json

@@ -0,0 +1,13 @@
+{
+  "name": "factorio-rcon-snippets-for-node",
+  "type": "module",
+  "version": "0.1.0",
+  "private": true,
+  "author": "LemonNekoGH",
+  "devDependencies": {
+    "@dotenvx/dotenvx": "^1.48.4",
+    "rcon-client": "^4.2.5",
+    "typescript": "^5.9.2",
+    "typescript-to-lua": "^1.31.4"
+  }
+}

packages/factorio-rcon-snippets-for-node/src/factorio_yolo_dataset_collector_v0.ts

@@ -0,0 +1,91 @@
+import { readFile } from 'node:fs/promises'
+import { env } from 'node:process'
+import { Rcon } from 'rcon-client'
+import { BuildMode, LuaLibImportKind, LuaTarget, transpileString } from 'typescript-to-lua'
+
+async function main() {
+  const rconPort = env.RCON_PORT
+  const rconHost = env.RCON_HOST
+  const rconPassword = env.RCON_PASSWORD
+
+  const seed = '2913363151'
+
+  if (!rconPort || !rconHost || !rconPassword) {
+    throw new Error('RCON_PORT, RCON_HOST, and RCON_PASSWORD must be set')
+  }
+
+  const codePath = new URL(
+    '../../factorio-rcon-snippets-for-vscode/src/factorio_yolo_dataset_collector_v0.ts',
+    import.meta.url,
+  )
+
+  const code = (await readFile(codePath)).toString()
+
+  const transpileResult = transpileString(
+    code.trim(),
+    {
+      luaTarget: LuaTarget.LuaJIT,
+      luaLibImport: LuaLibImportKind.Inline,
+      buildMode: BuildMode.Default,
+      noCheck: true,
+      noHeader: true,
+      noImplicitSelf: true,
+    },
+  )
+
+  if (transpileResult.diagnostics.length > 0) {
+    transpileResult.diagnostics.forEach((diagnostic) => {
+      if (typeof diagnostic.messageText === 'string') {
+        console.error(`Error: ${diagnostic.messageText} at ${diagnostic.start}`)
+      }
+      else {
+        console.error(`Error: ${diagnostic.messageText.messageText} at ${diagnostic.start}`)
+      }
+    })
+
+    return
+  }
+
+  if (!transpileResult.file || !transpileResult.file.lua) {
+    throw new Error('Transpilation did not produce valid Lua code.')
+  }
+
+  const entitiesNames = [
+    'assembling-machine-1',
+    'assembling-machine-2',
+    'assembling-machine-3',
+    'transport-belt',
+    'fast-transport-belt',
+    'express-transport-belt',
+  ]
+
+  const rconConnection = await Rcon.connect({
+    host: rconHost,
+    port: Number(rconPort),
+    password: rconPassword,
+  })
+
+  let scriptFirstRun = true
+  for (const index in entitiesNames) {
+    const name = entitiesNames[index]
+    await rconConnection.send(`/sc helpers.write_file('factorio_yolo_dataset_v0/classes.txt', '${name}\n', true)`)
+
+    let firstRun = true
+    for (let i = 0; i < 20; i++) {
+      const type = i > 18 ? 'test' : i > 14 ? 'val' : 'train'
+      // remove the last line 'return ____exports'
+      let codeWithParams = transpileResult.file.lua.toString()
+      codeWithParams = codeWithParams.replace(/return ____exports/, '')
+      codeWithParams += `\n\n____exports.main('${name}', ${index}, '${seed}', ${scriptFirstRun}, ${firstRun}, '${type}')`
+      const response = await rconConnection.send(`/sc ${codeWithParams}`)
+      // eslint-disable-next-line no-console
+      console.log(response)
+      scriptFirstRun = false
+      firstRun = false
+    }
+  }
+
+  await rconConnection.end()
+}
+
+main().then()

packages/factorio-rcon-snippets-for-node/tsconfig.json

@@ -0,0 +1,26 @@
+{
+  "compilerOptions": {
+    "target": "ESNext",
+    "lib": [
+      "ESNext"
+    ],
+    "moduleDetection": "auto",
+    "module": "ESNext",
+    "moduleResolution": "bundler",
+    "resolveJsonModule": true,
+    "strict": true,
+    "strictNullChecks": true,
+    "noImplicitAny": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "outDir": "./dist",
+    "esModuleInterop": true,
+    "forceConsistentCasingInFileNames": true,
+    "isolatedModules": true,
+    "verbatimModuleSyntax": true,
+    "skipLibCheck": true
+  },
+  "include": [
+    "src/**/*.ts"
+  ]
+}

packages/factorio-rcon-snippets-for-vscode/package.json

@@ -1,5 +1,5 @@
 {
-  "name": "autorio.ts",
+  "name": "factorio-rcon-snippets-for-vscode",
   "version": "0.1.0",
   "private": true,
   "author": "LemonNekoGH",

packages/factorio-rcon-snippets-for-vscode/src/factorio_yolo_dataset_collector_v0.ts

@@ -0,0 +1,130 @@
+/**
+ * If we create and destroy entities in the same tick,
+ * the screenshots will be the last state, so we need to run this script multiple times.
+ */
+import type { MapPositionStruct } from 'factorio:prototype'
+import type { BoundingBox } from 'factorio:runtime'
+
+rendering.clear()
+
+const surface_name = 'nauvis'
+const surface = game.surfaces[surface_name]
+const player = game.players[1]
+const player_position = player.position
+const capture_center: MapPositionStruct = {
+  x: math.floor(player_position.x),
+  y: math.floor(player_position.y),
+}
+const resolution = { x: 1280, y: 1280 }
+
+function clear_entities() {
+  const entities = surface.find_entities()
+  for (const e of entities) {
+    e.destroy()
+  }
+}
+
+// when zoom_level is 2 and resolution is 1280x1280, a picture contains 20x20 tiles
+// when zoom_level is 4 and resolution is 1280x1280, a picture contains 10x10 tiles
+// so we need to calculate the capture_box based on the zoom_level
+function calc_capture_box(
+  zoom_level: number,
+): BoundingBox {
+  const tile_size = 32 * zoom_level
+  const tile_count_x = resolution.x / tile_size
+  const offset_x = tile_count_x / 2
+  const tile_count_y = resolution.y / tile_size
+  const offset_y = tile_count_y / 2
+
+  return {
+    left_top: { x: capture_center.x - offset_x, y: capture_center.y - offset_y },
+    right_bottom: { x: capture_center.x + offset_x, y: capture_center.y + offset_y },
+  }
+}
+
+function selection_box_to_yolo_label(capture_box: BoundingBox, selection_box: BoundingBox) {
+  const capture_box_width = capture_box.right_bottom.x - capture_box.left_top.x
+  const capture_box_height = capture_box.right_bottom.y - capture_box.left_top.y
+
+  const center_x = ((selection_box.left_top.x + selection_box.right_bottom.x) / 2 - capture_box.left_top.x) / capture_box_width
+  const center_y = ((selection_box.left_top.y + selection_box.right_bottom.y) / 2 - capture_box.left_top.y) / capture_box_height
+  const width = (selection_box.right_bottom.x - selection_box.left_top.x) / capture_box_width
+  const height = (selection_box.right_bottom.y - selection_box.left_top.y) / capture_box_height
+
+  return `${center_x} ${center_y} ${width} ${height}`
+}
+
+export function main(
+  name: string,
+  index: number,
+  seed: string,
+  script_first_run: boolean,
+  label_first_run: boolean,
+  type: 'train' | 'val' | 'test',
+) {
+  if (script_first_run) {
+    helpers.write_file('factorio_yolo_dataset_v0/detect.yaml', `
+train: images/train
+val: images/val
+test: images/test
+
+names:
+`, true)
+  }
+
+  if (label_first_run) {
+    helpers.write_file('factorio_yolo_dataset_v0/detect.yaml', `  ${index}: ${name}\n`, true)
+  }
+
+  const entity_position = {
+    x: capture_center.x + math.random(10) - 5,
+    y: capture_center.y + math.random(10) - 5,
+  }
+
+  // const entity_position = {
+  //   x: capture_center.x,
+  //   y: capture_center.y,
+  // }
+
+  clear_entities()
+
+  const entity = surface.create_entity({
+    name,
+    position: entity_position,
+  })
+
+  if (!entity) {
+    error(`Failed to create entity ${name}`)
+  }
+
+  const selection_box = entity.selection_box
+  rcon.print(`selection box: ${selection_box.left_top.x}, ${selection_box.left_top.y}, ${selection_box.right_bottom.x}, ${selection_box.right_bottom.y}`)
+
+  const daytime = math.random()
+  // const daytime = 1
+  surface.daytime = daytime
+
+  const zoom = math.random(0.5, 4)
+  // const zoom = 2
+  player.zoom = zoom
+
+  const capture_box = calc_capture_box(zoom)
+
+  rcon.print(`capture box: ${capture_box.left_top.x}, ${capture_box.left_top.y}, ${capture_box.right_bottom.x}, ${capture_box.right_bottom.y}`)
+
+  // name standard: {name}_{daytime}_{zoom}_{surface}_{seed}_{left_top_x}_{left_top_y}_{right_bottom_x}_{right_bottom_y}
+  const name_with_label = `${name}_${daytime}_${zoom}_${surface_name}_${seed}_${selection_box.left_top.x}_${selection_box.left_top.y}_${selection_box.right_bottom.x}_${selection_box.right_bottom.y}`
+
+  game.take_screenshot({
+    position: capture_center,
+    resolution,
+    daytime,
+    zoom,
+    path: `factorio_yolo_dataset_v0/images/${type}/${name_with_label}.jpg`,
+  })
+
+  const yolo_label = selection_box_to_yolo_label(capture_box, selection_box)
+  rcon.print(`yolo label: ${index} ${yolo_label}`)
+
+  helpers.write_file(`factorio_yolo_dataset_v0/labels/${type}/${name_with_label}.txt`, `${index} ${yolo_label}`, true)
+}

packages/vscode-factorio-rcon-evaluator/src/codelens.ts

@@ -31,6 +31,12 @@ export function getFunctionContent(code: string, symbol: DocumentSymbol) {
 }
 
 export class EvaluatorCodeLensProvider implements CodeLensProvider {
+  #ctx: Context
+
+  constructor(ctx: Context) {
+    this.#ctx = ctx
+  }
+
   async provideCodeLenses(document: TextDocument, _: CancellationToken) {
     if (document.languageId !== 'typescript') {
       return []
@@ -54,8 +60,8 @@ export class EvaluatorCodeLensProvider implements CodeLensProvider {
   }
 }
 
-export function registerCodeLens(_: Context): Disposable[] {
+export function registerCodeLens(ctx: Context): Disposable[] {
   return [
-    languages.registerCodeLensProvider({ language: 'typescript' }, new EvaluatorCodeLensProvider()),
+    languages.registerCodeLensProvider({ language: 'typescript' }, new EvaluatorCodeLensProvider(ctx)),
   ]
 }

pnpm-lock.yaml

@@ -118,7 +118,22 @@ importers:
         specifier: ^1.31.4
         version: 1.31.4(typescript@5.9.2)
 
-  packages/factorio-rcon-snippets:
+  packages/factorio-rcon-snippets-for-node:
+    devDependencies:
+      '@dotenvx/dotenvx':
+        specifier: ^1.48.4
+        version: 1.48.4
+      rcon-client:
+        specifier: ^4.2.5
+        version: 4.2.5
+      typescript:
+        specifier: ^5.9.2
+        version: 5.9.2
+      typescript-to-lua:
+        specifier: ^1.31.4
+        version: 1.31.4(typescript@5.9.2)
+
+  packages/factorio-rcon-snippets-for-vscode:
     devDependencies:
       lua-types:
         specifier: ^2.13.1