diff --git a/backend/gpu/CLIP.py b/backend/gpu/CLIP.py new file mode 100644 index 0000000..bfbe3d4 --- /dev/null +++ b/backend/gpu/CLIP.py @@ -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} \ No newline at end of file diff --git a/backend/gpu/main.py b/backend/gpu/main.py new file mode 100644 index 0000000..f37119e --- /dev/null +++ b/backend/gpu/main.py @@ -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 \ No newline at end of file diff --git a/backend/gpu/request_generation.py b/backend/gpu/request_generation.py new file mode 100644 index 0000000..1213e12 --- /dev/null +++ b/backend/gpu/request_generation.py @@ -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) \ No newline at end of file