vumichien commited on
Commit
b96d288
Β·
1 Parent(s): f8685aa

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +13 -2
main.py CHANGED
@@ -1,3 +1,7 @@
 
 
 
 
1
  from ultralytics import YOLO
2
  from base64 import b64encode
3
  from speech_recognition import AudioFile, Recognizer
@@ -12,13 +16,18 @@ from fastapi.responses import StreamingResponse
12
  from utils import tts, read_image_file, pil_to_base64, get_hist
13
  from typing import Optional
14
  from huggingface_hub import hf_hub_download
15
- import time
16
 
17
  from io import BytesIO
18
  import zipfile
 
 
19
 
20
  model_path = hf_hub_download(repo_id="ultralyticsplus/yolov8s", filename='yolov8s.pt')
21
  model = YOLO(model_path)
 
 
 
22
  CLASS = model.model.names
23
  defaul_bot_voice = "γŠγ―γ„γ‚ˆγ†γ”γ–γ„γΎγ™"
24
  area_thres = 0.3
@@ -35,7 +44,9 @@ async def predict_api(
35
  file: UploadFile = File(...),
36
  last_seen: Union[UploadFile, None] = File(None)
37
  ):
 
38
  image = read_image_file(await file.read())
 
39
  results = model.predict(image, show=False)[0]
40
  image = read_image_as_pil(image)
41
  masks, boxes = results.masks, results.boxes
@@ -43,7 +54,7 @@ async def predict_api(
43
  most_close = 0
44
  out_img = None
45
  diff_value = 0.5
46
- total_time = time.time()
47
  start_time = time.time()
48
  if boxes is not None:
49
  for xyxy, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls):
 
1
+ import time
2
+ total_time = time.time()
3
+ start_time = time.time()
4
+
5
  from ultralytics import YOLO
6
  from base64 import b64encode
7
  from speech_recognition import AudioFile, Recognizer
 
16
  from utils import tts, read_image_file, pil_to_base64, get_hist
17
  from typing import Optional
18
  from huggingface_hub import hf_hub_download
19
+
20
 
21
  from io import BytesIO
22
  import zipfile
23
+ print("Import time", time.time() - start_time)
24
+ start_time = time.time()
25
 
26
  model_path = hf_hub_download(repo_id="ultralyticsplus/yolov8s", filename='yolov8s.pt')
27
  model = YOLO(model_path)
28
+
29
+ print("Load model time", time.time() - start_time)
30
+
31
  CLASS = model.model.names
32
  defaul_bot_voice = "γŠγ―γ„γ‚ˆγ†γ”γ–γ„γΎγ™"
33
  area_thres = 0.3
 
44
  file: UploadFile = File(...),
45
  last_seen: Union[UploadFile, None] = File(None)
46
  ):
47
+ start_time = time.time()
48
  image = read_image_file(await file.read())
49
+ print("Read image", time.time() - start_time)
50
  results = model.predict(image, show=False)[0]
51
  image = read_image_as_pil(image)
52
  masks, boxes = results.masks, results.boxes
 
54
  most_close = 0
55
  out_img = None
56
  diff_value = 0.5
57
+ # total_time = time.time()
58
  start_time = time.time()
59
  if boxes is not None:
60
  for xyxy, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls):