# Ensure the model directory exists
import os
import json
model_dir = os.path.join("model")
os.makedirs(model_dir, exist_ok=True)

classes = ['hotdog', 'not hotdog']
pretrained = True
input_size = [3, 256, 256]
normalize_mean = [0.5, 0.5, 0.5]
normalize_std = [0.5, 0.5, 0.5]

# Updated config data
config = {
    "model_name": "ResNet18",
    "num_classes": len(classes),
    "classes": classes,
    "input_size": input_size,  # Format: [channels, height, width]
    "normalize_mean": normalize_mean,  # Mean for normalization
    "normalize_std": normalize_std,    # Std for normalization
    "pretrained": pretrained           # Whether the model was pretrained
}

# Save the config
config_path = os.path.join(model_dir, "config.json")
with open(config_path, "w") as f:
    json.dump(config, f, indent=4)

print(f"Config saved to {config_path}")