Spaces:
Sleeping
Sleeping
Update video_processing.py
Browse files- 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 |
-
|
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 |
-
#
|
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 |
|