Update README.md

README.md

@@ -39,7 +39,6 @@ import torch
 
 # get device
 device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
-device = 'cpu'
 
 # load model
 # IF you cannot connect to huggingface.co, you can download the repo and set from_pretrained path as the loacl dir path, replacing "yuxindu/segvol"
@@ -109,3 +108,56 @@ print('done')
 ```
 
 ### Training script
+
+```python
+from transformers import AutoModel, AutoTokenizer
+import torch
+
+# get device
+device = torch.device("cuda:1" if torch.cuda.is_available() else "CPU")
+
+# load model
+# IF you cannot connect to huggingface.co, you can download the repo and set from_pretrained path as the loacl dir path, replacing "yuxindu/segvol"
+clip_tokenizer = AutoTokenizer.from_pretrained("./segvol")
+model = AutoModel.from_pretrained("./segvol", trust_remote_code=True, test_mode=False)
+model.model.text_encoder.tokenizer = clip_tokenizer
+model.train()
+model.to(device)
+print('model load done')
+
+# set case path
+# you can download this case from huggingface yuxindu/segvol files and versions
+ct_path = 'path/to/Case_image_00001_0000.nii.gz'
+gt_path = 'path/to/Case_label_00001.nii.gz'
+
+# set categories, corresponding to the unique values(1, 2, 3, 4, ...) in ground truth mask
+categories = ["liver", "kidney", "spleen", "pancreas"]
+
+# generate npy data format
+ct_npy, gt_npy = model.processor.preprocess_ct_gt(ct_path, gt_path, category=categories)
+# IF you have download our 25 processed datasets, you can skip to here with the processed ct_npy, gt_npy files
+
+# go through train transform
+data_item = model.processor.train_transform(ct_npy, gt_npy)
+
+# training example
+# add batch dim manually
+image, gt3D = data_item["image"].unsqueeze(0).to(device), data_item["label"].unsqueeze(0).to(device) # add batch dim
+
+loss_step_avg = 0
+for cls_idx in range(len(categories)):
+    # optimizer.zero_grad()
+    organs_cls = categories[cls_idx]
+    labels_cls = gt3D[:, cls_idx]
+    print(image.shape, organs_cls, labels_cls.shape)
+    loss = model.forward_train(image, train_organs=organs_cls, train_labels=labels_cls)
+    loss_step_avg += loss.item()
+    loss.backward()
+    # optimizer.step()
+
+loss_step_avg /= len(categories)
+print(f'AVG loss {loss_step_avg}')
+
+# save ckpt
+model.save_pretrained('./ckpt')
+```