Rename model.py to train.py

model.py → train.py

@@ -1,16 +1,15 @@
-import torch
-import torchvision.transforms as transforms
-from torchvision.datasets import ImageFolder
-from torchvision import models
 import os
 import tarfile
 import shutil
-from PIL import Image
+import torch
+import torchvision.models as models
+import torchvision.transforms as transforms
+from torchvision.datasets import ImageFolder
+
 
 # Set the paths
 data_dir = "C:/Users/kazia/Downloads/Trees"
 train_dir = os.path.join(data_dir, "Train")
-test_dir = os.path.join(data_dir, "Test")
 checkpoint_dir = "resnet_c_s_redwood_finetuned"
 tar_file = "resnet_c_s_redwood_finetuned.tar.gz"
 
@@ -77,6 +76,7 @@ torch.save(model.state_dict(), os.path.join(checkpoint_dir, "model.pth"))
 print("Creating tar.gz file...")
 with tarfile.open(tar_file, "w:gz") as tar:
     tar.add(checkpoint_dir, arcname=os.path.basename(checkpoint_dir))
+
 # Move the tar.gz file to the desired location
 destination_path = "C:/Users/kazia/Downloads/resnet_c_s_redwood_finetuned.tar.gz"
 shutil.move(tar_file, destination_path)
@@ -84,4 +84,4 @@ shutil.move(tar_file, destination_path)
 # Cleanup the intermediate checkpoint directory
 shutil.rmtree(checkpoint_dir)
 
-print("Fine-tuned model saved and compressed as tar.gz.")
+print("Fine-tuned model saved and compressed as tar.gz.")