Add files via upload

This commit is contained in:
John Kang 2022-09-23 18:22:30 -06:00 committed by GitHub
parent 111687b015
commit b04c0243a5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 54 additions and 0 deletions

22
backend/gpu/CLIP.py Normal file
View File

@ -0,0 +1,22 @@
import clip
import torch
import numpy as np
import cv2
from PIL import Image
import json
def embed(frames,model,preprocess,device):
x = json.loads(frames)
frames = np.array(x,dtype=np.uint8)
print(frames.shape)
frames_pil = [Image.fromarray(im_cv) for im_cv in frames]
frames_preprocess = [preprocess(frame) for frame in frames_pil]
frames_preprocess = torch.stack(frames_preprocess)
print(frames_preprocess.shape)
with torch.no_grad():
vid_encodings = model.encode_image(frames_preprocess.to(device).half())
print(vid_encodings.shape)
vid_encodings_json = json.dumps(vid_encodings.numpy().tolist())
return {'clip_encodings': vid_encodings_json}

20
backend/gpu/main.py Normal file
View File

@ -0,0 +1,20 @@
from CLIP import embed
import clip
import torch
import numpy as np
import cv2
from PIL import Image
import json
from fastapi import FastAPI
from typing import Dict
app = FastAPI()
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)
@app.post('/')
async def handle(request: Dict):
frames = request['video_frames']
response = embed(frames, model, preprocess, device)
return response

View File

@ -0,0 +1,12 @@
import requests
import json
import numpy as np
num_frames = 1
shape = (num_frames, 300, 300, 3)
random_frames = np.random.randint(0,255,size=shape,dtype = np.uint8)
fr_json = json.dumps(random_frames.tolist())
r = requests.post('http://127.0.0.1:8000/', json={"video_frames": fr_json})
print(r.content)