foryahasake commited on
Commit
26efb89
·
verified ·
1 Parent(s): b6fc81f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -1
app.py CHANGED
@@ -36,8 +36,12 @@ my_metadata = MetadataCatalog.get("dbmdz_coco_all")
36
  my_metadata.thing_classes = ["None", "BAD_BILLBOARD","BROKEN_SIGNAGE","CLUTTER_SIDEWALK","CONSTRUCTION_ROAD","FADED_SIGNAGE","GARBAGE","GRAFFITI","POTHOLES","SAND_ON_ROAD","UNKEPT_FACADE"]
37
 
38
 
 
 
 
39
 
40
  def predict_frame(frame,_):
 
41
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3
42
  predictor = DefaultPredictor(cfg)
43
  outputs = predictor(frame)
@@ -45,6 +49,46 @@ def predict_frame(frame,_):
45
  out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
46
  return out.get_image()
47
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
  @spaces.GPU
49
  def inference(image_url, image, min_score):
50
  if not torch.cuda.is_available():
@@ -74,7 +118,7 @@ def inference(image_url, image, min_score):
74
 
75
  @spaces.GPU
76
  def process_vid(video_path):
77
-
78
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3
79
  if not torch.cuda.is_available():
80
  cfg.MODEL.DEVICE = "cpu"
 
36
  my_metadata.thing_classes = ["None", "BAD_BILLBOARD","BROKEN_SIGNAGE","CLUTTER_SIDEWALK","CONSTRUCTION_ROAD","FADED_SIGNAGE","GARBAGE","GRAFFITI","POTHOLES","SAND_ON_ROAD","UNKEPT_FACADE"]
37
 
38
 
39
+
40
+
41
+
42
 
43
  def predict_frame(frame,_):
44
+
45
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3
46
  predictor = DefaultPredictor(cfg)
47
  outputs = predictor(frame)
 
49
  out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
50
  return out.get_image()
51
 
52
+ @spaces.GPU
53
+ def opt_process_vid(video_path):
54
+
55
+
56
+ cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3
57
+ cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
58
+ predictor = DefaultPredictor(cfg)
59
+ v = VideoVisualizer(my_metadata,ColorMode.IMAGE)
60
+ cap = cv2.VideoCapture(video_path)
61
+ frame_size = (int(cap.get(3)), int(cap.get(4)))
62
+ fps = int(cap.get(5))
63
+ vid_fourcc= int(cap.get(cv2.CAP_PROP_FOURCC))
64
+ output_path = '/content/drive/MyDrive/ColabNotebooks/gradio-exp/output.mp4'
65
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
66
+ video_writer = cv2.VideoWriter(output_path,fourcc, fps, frame_size)
67
+ num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
68
+ skipped_viz = 5
69
+ skipped_outputs= 15
70
+ # Process the video
71
+ for i in tqdm.tqdm(range(num_frames)):
72
+ ret, frame = cap.read()
73
+ if not ret:
74
+ break
75
+ if i % skipped_outputs == 0:
76
+ # Get prediction results for this frame
77
+ outputs = predictor(frame)
78
+ if i % skipped_viz == 0:
79
+ # Draw a visualization of the predictions
80
+ frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
81
+ visualization = v.draw_instance_predictions(frame, outputs["instances"].to("cpu"))
82
+ visualization = cv2.cvtColor(visualization.get_image(), cv2.COLOR_RGB2BGR)
83
+ video_writer.write(visualization)
84
+ # Release resources
85
+ cap.release()
86
+ video_writer.release()
87
+ torch.cuda.empty_cache()
88
+
89
+ return output_path
90
+
91
+
92
  @spaces.GPU
93
  def inference(image_url, image, min_score):
94
  if not torch.cuda.is_available():
 
118
 
119
  @spaces.GPU
120
  def process_vid(video_path):
121
+
122
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3
123
  if not torch.cuda.is_available():
124
  cfg.MODEL.DEVICE = "cpu"