Update README.md

README.md

@@ -53,7 +53,7 @@ Video classification model
    # Enable multi-GPU support
    net = torch.nn.DataParallel(net)
    torch.backends.cudnn.benchmark = True
-   state = torch.load(model_path, map_location=torch.device('cuda'))
+   state = torch.load(model_path, map_location=torch.device('cpu'))
    net.load_state_dict(state['net'])
    net.eval()
 
@@ -63,7 +63,61 @@ Video classification model
    img = img.resize((224, 224))
    img_tensor = torch.tensor(np.array(img)).unsqueeze(0).to('cuda')
 
-   # Extract features from the image using the ResNet50 model
+   # Extract features from the image
+   outputs = net(img_tensor)
+   ```
+
+Frame classification model
+
+   ```python
+   import torch
+   from PIL import Image
+   from model_loader import build_model
+
+   # Load the model
+   net = build_model(mode='classify')
+   model_path = 'Frame classification models'
+
+   # Enable multi-GPU support
+   net = torch.nn.DataParallel(net)
+   torch.backends.cudnn.benchmark = True
+   state = torch.load(model_path, map_location=torch.device('cpu'))
+   net.load_state_dict(state['net'])
+   net.eval()
+
+   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')
+
+   # Extract features from the image
+   outputs = net(img_tensor)
+   ```
+
+Non-surgical object detection model
+
+   ```python
+   import torch
+   from PIL import Image
+   from model_loader import build_model
+
+   # Load the model
+   net = build_model(mode='mask')
+   model_path = 'Frame classification models'
+
+   # Enable multi-GPU support
+   net = torch.nn.DataParallel(net)
+   torch.backends.cudnn.benchmark = True
+   state = torch.load(model_path, map_location=torch.device('cpu'))
+   net.load_state_dict(state['net'])
+   net.eval()
+
+   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')
+
+   # Extract features from the image
    outputs = net(img_tensor)
    ```