Update README.md

README.md

@@ -69,6 +69,7 @@ Usage
 **Video classification models** are employed in the step **2** of the data curation pipeline to classify a video storyboard as either surgical or non-surgical, the models usage is as follows:
    ```python
    import torch
+   import torchvision
    from PIL import Image
    from model_loader import build_model
 
@@ -87,7 +88,14 @@ Usage
    img_path = 'path/to/your/image.jpg'
    img = Image.open(img_path)
    img = img.resize((224, 224))
-   img_tensor = torch.tensor(np.array(img)).unsqueeze(0).to('cuda')
+   transform = torchvision.transforms.Compose([
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(
+        (0.4299694, 0.29676908, 0.27707579), 
+        (0.24373249, 0.20208984, 0.19319402)
+    )
+   ])
+   img_tensor = transform(img).unsqueeze(0).to('cuda')
 
    # Extract features from the image
    outputs = net(img_tensor)
@@ -114,7 +122,14 @@ Usage
    img_path = 'path/to/your/image.jpg'
    img = Image.open(img_path)
    img = img.resize((224, 224))
-   img_tensor = torch.tensor(np.array(img)).unsqueeze(0).to('cuda')
+   transform = torchvision.transforms.Compose([
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(
+        (0.4299694, 0.29676908, 0.27707579), 
+        (0.24373249, 0.20208984, 0.19319402)
+    )
+   ])
+   img_tensor = transform(img).unsqueeze(0).to('cuda')
 
    # Extract features from the image
    outputs = net(img_tensor)
@@ -141,7 +156,14 @@ Usage
    img_path = 'path/to/your/image.jpg'
    img = Image.open(img_path)
    img = img.resize((224, 224))
-   img_tensor = torch.tensor(np.array(img)).unsqueeze(0).to('cuda')
+   transform = torchvision.transforms.Compose([
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(
+        (0.4299694, 0.29676908, 0.27707579), 
+        (0.24373249, 0.20208984, 0.19319402)
+    )
+   ])
+   img_tensor = transform(img).unsqueeze(0).to('cuda')
 
    # Extract features from the image
    outputs = net(img_tensor)