scfive commited on
Commit
2a43d99
·
verified ·
1 Parent(s): 61c4868

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -71
app.py CHANGED
@@ -1,77 +1,25 @@
 
 
1
  import gradio as gr
2
  import os
3
- import tempfile
4
- import cv2
5
- import numpy as np
6
- from mmdet.apis import DetInferencer
7
 
8
- # Helper to load model
9
- inferencer = None
10
- def load_model(config_path, checkpoint_path):
11
- global inferencer
12
- inferencer = DetInferencer(model=config_path, weights=checkpoint_path)
13
- return "Model loaded."
14
 
15
- def infer_image(image):
16
- if inferencer is None:
17
- return "Please load a model first.", None
18
- result = inferencer(image)
19
- vis = result["visualization"]
20
- if isinstance(vis, list):
21
- vis = vis[0]
22
- return "", vis
23
-
24
- def infer_video(video):
25
- if inferencer is None:
26
- return "Please load a model first.", None
27
- temp_dir = tempfile.mkdtemp()
28
- cap = cv2.VideoCapture(video)
29
- fps = cap.get(cv2.CAP_PROP_FPS)
30
- w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
31
- h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
32
- out_path = os.path.join(temp_dir, "result.mp4")
33
- fourcc = cv2.VideoWriter_fourcc(*"mp4v")
34
- out = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
35
- while True:
36
- ret, frame = cap.read()
37
- if not ret:
38
- break
39
- result = inferencer(frame)
40
- vis = result["visualization"]
41
- if isinstance(vis, list):
42
- vis = vis[0]
43
- out.write(vis[:,:,::-1])
44
- cap.release()
45
- out.release()
46
- return "", out_path
47
-
48
- def ui():
49
- with gr.Blocks() as demo:
50
- gr.Markdown("# SpecDETR Demo: Image and Video Detection\nUpload your config (.py) and checkpoint (.pth) to start.")
51
- with gr.Row():
52
- config = gr.File(label="Config File (.py)")
53
- checkpoint = gr.File(label="Checkpoint (.pth)")
54
- load_btn = gr.Button("Load Model")
55
- load_status = gr.Textbox(label="Status", interactive=False)
56
- load_btn.click(load_model, inputs=[config, checkpoint], outputs=load_status)
57
- with gr.Tab("Image"):
58
- img_input = gr.Image(type="numpy")
59
- img_output = gr.Image()
60
- img_btn = gr.Button("Detect on Image")
61
- img_status = gr.Textbox(label="Status", interactive=False)
62
- img_btn.click(infer_image, inputs=img_input, outputs=[img_status, img_output])
63
- with gr.Tab("Video"):
64
- vid_input = gr.Video()
65
- vid_output = gr.Video()
66
- vid_btn = gr.Button("Detect on Video")
67
- vid_status = gr.Textbox(label="Status", interactive=False)
68
- vid_btn.click(infer_video, inputs=vid_input, outputs=[vid_status, vid_output])
69
- return demo
70
-
71
- demo = ui()
72
-
73
- def main():
74
- demo.launch()
75
 
76
  if __name__ == "__main__":
77
- main()
 
1
+ import sys, os
2
+ sys.path.append(os.path.abspath(os.path.dirname(__file__)))
3
  import gradio as gr
4
  import os
5
+ from inference_custom import main as samwise_infer
 
 
 
6
 
7
+ def inference(video, prompt):
8
+ output_path = "output_segmented.mp4"
9
+ # 'video' is a file path string provided by Gradio
10
+ samwise_infer(video, prompt, output_path, "models/samwise.pth")
11
+ return output_path
 
12
 
13
+ demo = gr.Interface(
14
+ fn=inference,
15
+ inputs=[
16
+ gr.Video(label="Upload Video"),
17
+ gr.Textbox(label="Text Prompt", placeholder="Describe what to segment")
18
+ ],
19
+ outputs=gr.Video(label="Segmented Output"),
20
+ title="SAMWISE Video Segmentation",
21
+ description="Upload a video and enter a prompt to segment objects with SAMWISE."
22
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23
 
24
  if __name__ == "__main__":
25
+ demo.launch()