khengkok commited on
Commit
11fdcc2
·
1 Parent(s): 9d6af19
Files changed (2) hide show
  1. app.py +42 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ultralytics import YOLO
2
+ from PIL import Image
3
+ import gradio as gr
4
+ from huggingface_hub import snapshot_download
5
+ from huggingface_hub import hf_hub_download
6
+ import os
7
+
8
+ MODEL_ID = "goldfish_v2.pt"
9
+ REPO_ID = "ITI121-25S2/112233A"
10
+
11
+ def load_model(repo_id):
12
+ ## The following code shows how to download directory from Hugging Face model
13
+ # download_dir = snapshot_download(repo_id)
14
+ # print(download_dir)
15
+ # path = os.path.join(download_dir, "best_int8_openvino_mode")
16
+
17
+ # Download the single model file from Hugging Face model repository
18
+
19
+ path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_ID)
20
+
21
+ print(path)
22
+ detection_model = YOLO(path, task='detect')
23
+ return detection_model
24
+
25
+
26
+ def predict(pilimg):
27
+
28
+ source = pilimg
29
+ # x = np.asarray(pilimg)
30
+ # print(x.shape)
31
+ result = detection_model.predict(source, conf=0.5, iou=0.6)
32
+ img_bgr = result[0].plot()
33
+ out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # RGB-order PIL image
34
+
35
+ return out_pilimg
36
+
37
+ detection_model = load_model(REPO_ID)
38
+
39
+ gr.Interface(fn=predict,
40
+ inputs=gr.Image(type="pil"),
41
+ outputs=gr.Image(type="pil")
42
+ ).launch(share=True)
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ ultralytics
2
+ huggingface_hub
3
+ gradio