hyfeng017 commited on
Commit
2decca9
·
verified ·
1 Parent(s): c14603a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ultralytics import YOLO
2
+ from PIL import Image
3
+ import gradio as gr
4
+ from huggingface_hub import snapshot_download
5
+ import os
6
+
7
+ model_path = "/Users/markk/Downloads/best_int8_openvino_model"
8
+
9
+ def load_model(repo_id):
10
+ download_dir = snapshot_download(repo_id)
11
+ print(download_dir)
12
+ path = os.path.join(download_dir, "best_int8_openvino_model")
13
+ print(path)
14
+ detection_model = YOLO(path, task='detect')
15
+ return detection_model
16
+
17
+
18
+ def predict(pilimg):
19
+
20
+ source = pilimg
21
+ # x = np.asarray(pilimg)
22
+ # print(x.shape)
23
+ result = detection_model.predict(source, conf=0.5, iou=0.6)
24
+ img_bgr = result[0].plot()
25
+ out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # RGB-order PIL image
26
+
27
+ return out_pilimg
28
+
29
+
30
+ REPO_ID = "khengkok/balloon_yolov8"
31
+ detection_model = load_model(REPO_ID)
32
+
33
+ gr.Interface(fn=predict,
34
+ inputs=gr.Image(type="pil"),
35
+ outputs=gr.Image(type="pil")
36
+ ).launch(share=True)