Update README.md

README.md

@@ -67,8 +67,15 @@ sample = {
 """
 sample = dataset[index]
 
+depth: Image.Image = sample["depth"]
+rgb: Image.Image = sample["rgb"]
+gt: Image.Image = sample["gt"]
+name: str = sample["name"]
+
 
 """
+1. Preprocessing step
+
 preprocessed_sample = {
     'rgb': tensor([[[[-0.8507, ....0365]]]]), 
     'gt': tensor([[[[0., 0., 0...., 0.]]]]), 
@@ -78,12 +85,17 @@ preprocessed_sample = {
 preprocessed_sample = image_processor.preprocess(sample)
 
 """
+2. Prediction step
+
 output = {
     'logits': tensor([[[[-5.1966, ...ackward0>)
 }
 """
 output = model(preprocessed_sample["rgb"], preprocessed_sample["depth"])
 
+"""
+3. Postprocessing step
+"""
 postprocessed_sample: np.ndarray = image_processor.postprocess(
     output["logits"], [sample["gt"].size[1], sample["gt"].size[0]]
 )
@@ -93,7 +105,8 @@ prediction = Image.fromarray(np.uint8(cm.gist_earth(postprocessed_sample) * 255)
 Show the predicted salient map and the corresponding ground-truth(GT)
 """
 prediction.show()
-sample["gt"].show()
+gt.show()
+
 ```
 
 ### Downstream Use [optional]