pytorch reads frames from video for image classification, all predicting the same category
I have 3 categories of fish images, each category has 1000 images, when I trained the model, the classification accuracy was 97%, and when I used the folder images for prediction, the accuracy was no problem.
But when I replace the picture with a video, and cut out each frame from the video for image classification, no matter which category of video it is, all frames are predicted to be category 1:&quot;stingray&quot;.Why?
#!/usr/bin/env python
# coding: utf-8

import torch
from torchvision import transforms
import torchvision.models as models
import cv2
import torch.nn.functional as F


CLASSES = {0:&quot;goldfish&quot;, 1:&quot;stingray&quot;, 2:&quot;tench&quot;}
BATCH_SIZE = 4
IMG_SIZE = (400, 400)
TRANSFORM_IMG = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize(IMG_SIZE),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225] )
    ])

# model
device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
model = models.vgg19_bn(pretrained=False, num_classes=3)
model.to(device)
model.load_state_dict(torch.load('checkpoint.pt'))
model.eval()


videoCapture = cv2.VideoCapture(r'D:/video/Goldfish.mp4')
fps = videoCapture.get(cv2.CAP_PROP_FPS)
size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT)))

ps = 25
fourcc = cv2.VideoWriter_fourcc(*'DIVX')
videoWriter = cv2.VideoWriter(&quot;D:/goldfish.mp4&quot;, fourcc, fps, size)

with torch.no_grad():
    success, frame = videoCapture.read()
    while success:
        # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image_tensor = TRANSFORM_IMG(frame)
        image_tensor = image_tensor.unsqueeze(0) 
        test_input = image_tensor.to(device)
        outputs = model(test_input)
        _, predicted = torch.max(outputs, 1)
        probability =  F.softmax(outputs, dim=1)
        top_probability, top_class = probability.topk(1, dim=1)
        predicted = predicted.cpu().detach().numpy()
        predicted = predicted.tolist()[0]
        label = CLASSES[predicted]
        top_probability = top_probability.cpu().detach().numpy()
        top_probability = top_probability.tolist()[0][0]
        top_probability = '%.2f%%' % (top_probability * 100)
        print(top_probability)
        print(label) #all the label is stingray############################################
        frame = cv2.putText(frame, label+': '+top_probability, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 2)
        videoWriter.write(frame)
        success, frame = videoCapture.read()
    videoWriter.release()


When I use this code,there is no problem.I convert image channel from BGR to RGB for pytorch.
import torch
from torchvision import transforms
import torchvision.models as models
import cv2
import torch.nn.functional as F
import copy

CLASSES = {0:&quot;goldfish&quot;, 1:&quot;stingray&quot;, 2:&quot;tench&quot;}
BATCH_SIZE = 4
IMG_SIZE = (400, 400)
TRANSFORM_IMG = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize(IMG_SIZE),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225] )
    ])

# model
device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
model = models.vgg19_bn(pretrained=False, num_classes=3)
model.to(device)
model.load_state_dict(torch.load('checkpoint.pt'))
model.eval()


videoCapture = cv2.VideoCapture(r'video/Goldfish.mp4')
fps = videoCapture.get(cv2.CAP_PROP_FPS)
size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT)))

ps = 25
fourcc = cv2.VideoWriter_fourcc(*'DIVX')
videoWriter = cv2.VideoWriter(r&quot;D:/goldfish.mp4&quot;, fourcc, fps, size)

with torch.no_grad():
    success, frame = videoCapture.read()
    while success:
        frame_copy = copy.deepcopy(frame) 
        frame_copy = cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
        image_tensor = TRANSFORM_IMG(frame_copy)
        image_tensor = image_tensor.unsqueeze(0) 
        test_input = image_tensor.to(device)
        outputs = model(test_input)
        _, predicted = torch.max(outputs, 1)
        probability =  F.softmax(outputs, dim=1)
        top_probability, top_class = probability.topk(1, dim=1)
        predicted = predicted.cpu().detach().numpy()
        predicted = predicted.tolist()[0]
        label = CLASSES[predicted]
        top_probability = top_probability.cpu().detach().numpy()
        top_probability = top_probability.tolist()[0][0]
        top_probability = '%.2f%%' % (top_probability * 100)
        print(top_probability)
        print(label)
        frame = cv2.putText(frame, label+': '+top_probability, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 2)
        videoWriter.write(frame)
        success, frame = videoCapture.read()
    videoWriter.release()

