VoidCodes commited on
Commit
3ae786e
·
1 Parent(s): 28ac872

refactor: replace greeting function with YOLO model integration and Gradio interface for image prediction

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