mtg-chair-ai/generate-card-captions.js

94 lines
2.7 KiB
JavaScript

//TODO:
// Error handling for bad args
// Error handling for read/write errors
// Retry on failed requests (toggle with flag?)
const args = {
modelVersion: "latest", // -m --model-version
apiVersion: "2023-10-01", // -a --api-version
input: "unique-artwork-20240519090229.json", // -i --input
output: "output.json", // -o --output
}
const cliArgs = process.argv.slice(2)
while (cliArgs.length) {
const arg = cliArgs.shift()
if (arg === "-m" || arg === "--model-version") {
args.modelVersion = cliArgs.shift()
continue
}
if (arg === "-a" || arg === "--api-version") {
args.apiVersion = cliArgs.shift()
continue
}
if (arg === "-i" || arg === "--input") {
args.input = cliArgs.shift()
continue
}
if (arg === "-o" || arg === "--output") {
args.output = cliArgs.shift()
}
}
console.log("Running with args: ", args)
const fs = require("fs")
var cardArray = JSON.parse(fs.readFileSync(args.input, "utf8"))
console.log(`Read ${cardArray.length} cards from ${args.input}`)
const largeImages = cardArray.filter((card) => card.image_uris?.large)
let simpleCardArray = largeImages.map((card) => {
return {
imageUrl: card.image_uris.large,
name: card.name,
id: card.id,
oracleId: card.oracle_id,
url: card.scryfall_uri,
mediumImageUrl:
card.image_uris.normal || card.image_uris.small || card.image_uris.large,
caption: null,
}
})
console.log(`Filtered ${simpleCardArray.length} cards with large images`)
const getCaptionForImageUrl = async (url) => {
try {
const response = await fetch(
`${AZURE_VISION_RESOURCE_URL}/computervision/imageanalysis:analyze?features=denseCaptions&model-version=${modelVersion}&language=en&gender-neutral-caption=false&api-version=${apiVersion}`,
{
method: "POST",
headers: {
"Content-Type": "application/json",
"Ocp-Apim-Subscription-Key":
process.env.AZURE_VISION_SUBSCRIPTION_KEY,
},
body: JSON.stringify({
url,
}),
}
)
const data = await response.json()
return data.denseCaptionsResult.values
.map((caption) => caption.text)
.join("|")
} catch (error) {
console.error("Error fetching caption for card: ", url)
console.error(error)
return null
}
}
const addCaptionsToCards = async (cards) => {
fs.writeFileSync(args.output, "[", "utf8")
for (const card of cards) {
const caption = await getCaptionForImageUrl(card.imageUrl)
console.log(`Card ${card.id} caption: ${caption}`)
card.caption = caption
fs.appendFileSync(args.output, JSON.stringify(card), "utf8")
fs.appendFileSync(args.output, ",\n", "utf8")
}
fs.appendFileSync(args.output, "]", "utf8")
return cards
}
addCaptionsToCards(simpleCardArray)