efraim1011 commited on
Commit
9ff5c58
1 Parent(s): 2ed29c0

feature 'real time' added

Browse files
Files changed (1) hide show
  1. app.py +44 -2
app.py CHANGED
@@ -210,6 +210,43 @@ def process_video(
210
  sink.write_frame(frame)
211
  return result_file_path
212
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
213
 
214
  confidence_threshold_component = gr.Slider(
215
  minimum=0,
@@ -265,7 +302,7 @@ with_class_agnostic_nms_component = gr.Checkbox(
265
 
266
  with gr.Blocks() as demo:
267
  gr.Markdown(MARKDOWN)
268
- with gr.Accordion("Confiduração", open=False):
269
  confidence_threshold_component.render()
270
  iou_threshold_component.render()
271
  with gr.Row():
@@ -340,6 +377,10 @@ with gr.Blocks() as demo:
340
  # ],
341
  # outputs=output_image_component
342
  # )
 
 
 
 
343
 
344
  image_submit_button_component.click(
345
  fn=process_image,
@@ -367,5 +408,6 @@ with gr.Blocks() as demo:
367
  ],
368
  outputs=output_video_component
369
  )
 
370
 
371
- demo.launch(debug=False, show_error=True, max_threads=1)
 
210
  sink.write_frame(frame)
211
  return result_file_path
212
 
213
+ def process_image_real_time(
214
+ input_image: np.ndarray,
215
+ categories: str = "safety helmet, safety glasses, hearing protectors, protective masks, safety shoes, protective gloves, seat belt, person, bicycle, car, motorcycle, airplane, dining table, bus, train, truck, boat, traffic light, fire hydrant, stop sign, parking meter, bench, bird, cat, dog, horse, sheep, cow, elephant, bear, zebra, giraffe, backpack, umbrella, handbag, tie, suitcase, frisbee, skis, snowboard, sports ball, kite, baseball bat, baseball glove, skateboard, surfboard, tennis racket, bottle, wine glass, cup, fork, knife, spoon, bowl, banana, apple, sandwich, orange, broccoli, carrot, hot dog, pizza, donut, cake, chair, couch, potted plant, bed, toilet, tv, laptop, mouse, remote, keyboard, cell phone, microwave, oven, toaster, sink, refrigerator, book, clock, vase, scissors, teddy bear, hair drier, toothbrush",
216
+ confidence_threshold: float = 0.3,
217
+ iou_threshold: float = 0.5,
218
+ # with_segmentation: bool = True,
219
+ with_confidence: bool = False,
220
+ with_class_agnostic_nms: bool = False,) -> np.ndarray:
221
+ # cleanup of old video files
222
+
223
+ remove_files_older_than(RESULTS, 30)
224
+
225
+ categories = process_categories(categories)
226
+ YOLO_WORLD_MODEL.set_classes(categories)
227
+ results = YOLO_WORLD_MODEL.infer(input_image, confidence=0.02)
228
+ detections = sv.Detections.from_inference(results)
229
+ detections = detections.with_nms(
230
+ class_agnostic=with_class_agnostic_nms,
231
+ threshold=iou_threshold
232
+ )
233
+ # if with_segmentation:
234
+ # detections.mask = inference_with_boxes(
235
+ # image=input_image,
236
+ # xyxy=detections.xyxy,
237
+ # model=EFFICIENT_SAM_MODEL,
238
+ # device=DEVICE
239
+ # )
240
+ output_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
241
+ output_image = annotate_image(
242
+ input_image=output_image,
243
+ detections=detections,
244
+ categories=categories,
245
+ with_confidence=with_confidence
246
+ )
247
+ return cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)
248
+
249
+
250
 
251
  confidence_threshold_component = gr.Slider(
252
  minimum=0,
 
302
 
303
  with gr.Blocks() as demo:
304
  gr.Markdown(MARKDOWN)
305
+ with gr.Accordion("Configuração", open=False):
306
  confidence_threshold_component.render()
307
  iou_threshold_component.render()
308
  with gr.Row():
 
377
  # ],
378
  # outputs=output_image_component
379
  # )
380
+ with gr.Tab(label="Tempo Real"):
381
+ with gr.Row():
382
+ input_real_time_component = gr.Image(label="Entrada", sources=["webcam"], streaming=True)
383
+ output_real_time_component = gr.Image(label="Saída")
384
 
385
  image_submit_button_component.click(
386
  fn=process_image,
 
408
  ],
409
  outputs=output_video_component
410
  )
411
+ input_real_time_component.change(fn=process_image_real_time,inputs=input_real_time_component, outputs=output_real_time_component)
412
 
413
+ demo.launch(debug=False, show_error=True, max_threads=1, share=True)