adirik commited on
Commit
95255f1
1 Parent(s): bc5dfe0

use gpu is available

Browse files
Files changed (3) hide show
  1. .DS_Store +0 -0
  2. app.py +10 -3
  3. requirements.txt +0 -1
.DS_Store DELETED
Binary file (6.15 kB)
 
app.py CHANGED
@@ -2,16 +2,23 @@ import torch
2
  import cv2
3
  import gradio as gr
4
  import numpy as np
5
- from PIL import Image, ImageDraw, ImageFont
6
  from transformers import OwlViTProcessor, OwlViTForObjectDetection
7
 
8
- model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32").eval()
 
 
 
 
 
 
 
 
9
  processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
10
 
11
 
12
  def query_image(img, text_queries):
13
  text_queries = text_queries.split(",")
14
- inputs = processor(text=text_queries, images=img, return_tensors="pt")
15
 
16
  with torch.no_grad():
17
  outputs = model(**inputs)
 
2
  import cv2
3
  import gradio as gr
4
  import numpy as np
 
5
  from transformers import OwlViTProcessor, OwlViTForObjectDetection
6
 
7
+
8
+ # Use GPU if available
9
+ if torch.cuda.is_available():
10
+ device = torch.device("cuda")
11
+ else:
12
+ device = torch.device("cpu")
13
+
14
+ model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32").to(device)
15
+ model.eval()
16
  processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
17
 
18
 
19
  def query_image(img, text_queries):
20
  text_queries = text_queries.split(",")
21
+ inputs = processor(text=text_queries, images=img, return_tensors="pt").to(device)
22
 
23
  with torch.no_grad():
24
  outputs = model(**inputs)
requirements.txt CHANGED
@@ -3,6 +3,5 @@
3
  numpy>=1.18.5
4
  torch>=1.7.0
5
  torchvision>=0.8.1
6
- Pillow
7
  transformers
8
  opencv-python
 
3
  numpy>=1.18.5
4
  torch>=1.7.0
5
  torchvision>=0.8.1
 
6
  transformers
7
  opencv-python