Spaces:
Sleeping
Sleeping
# 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}") | |