jschwab21 commited on
Commit
c0981f1
·
verified ·
1 Parent(s): a3a38fc

Update video_processing.py

Browse files
Files changed (1) hide show
  1. video_processing.py +3 -11
video_processing.py CHANGED
@@ -18,15 +18,10 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18
  model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
19
  processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
20
 
 
21
 
22
- def classify_frame(frame):
23
- categories = ["Joy", "Trust", "Fear", "Surprise", "Sadness", "Disgust", "Anger", "Anticipation"]
24
-
25
- # Load ResNet-50 model
26
- resnet50 = models.resnet50(pretrained=True)
27
- resnet50.eval().to(device)
28
 
29
- # Preprocess the image
30
  preprocess = transforms.Compose([
31
  transforms.Resize(256),
32
  transforms.CenterCrop(224),
@@ -36,15 +31,12 @@ def classify_frame(frame):
36
  input_tensor = preprocess(Image.fromarray(frame))
37
  input_batch = input_tensor.unsqueeze(0).to(device)
38
 
39
- # Predict with ResNet-50
40
  with torch.no_grad():
41
  output = resnet50(input_batch)
42
  probabilities = F.softmax(output[0], dim=0)
43
 
44
- # Create a numpy array from the probabilities of the categories
45
- # This example assumes each category is mapped to a model output directly
46
  results_array = np.array([probabilities[i].item() for i in range(len(categories))])
47
-
48
  return results_array
49
 
50
 
 
18
  model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
19
  processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
20
 
21
+ resnet50 = models.resnet50(pretrained=True).eval().to(device)
22
 
 
 
 
 
 
 
23
 
24
+ def classify_frame(frame):
25
  preprocess = transforms.Compose([
26
  transforms.Resize(256),
27
  transforms.CenterCrop(224),
 
31
  input_tensor = preprocess(Image.fromarray(frame))
32
  input_batch = input_tensor.unsqueeze(0).to(device)
33
 
34
+ # Use the globally loaded ResNet-50 model
35
  with torch.no_grad():
36
  output = resnet50(input_batch)
37
  probabilities = F.softmax(output[0], dim=0)
38
 
 
 
39
  results_array = np.array([probabilities[i].item() for i in range(len(categories))])
 
40
  return results_array
41
 
42