hotdog / src /config.py
asidfactory's picture
Uploaded coded initial version
2a0bba9 verified
raw
history blame contribute delete
878 Bytes
# 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}")