Update app.py

app.py

@@ -3,7 +3,9 @@ import time
 import cv2 # opencv2 package for python.
 import torch
 from pytube import YouTube
-
+from PIL import Image
+import numpy as np
+from transformers import pipeline
 
 segmentor = pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic")
 
@@ -22,6 +24,7 @@ def load(URL):
     frame_fps = (player.get(cv2.CAP_PROP_FPS))
     tog = 0
     return vid_cap,frame_num,frame_count,frame_fps,tog
+    
 
 def fw_fn(cur,last):
     next = cur+1
@@ -43,15 +46,33 @@ def pl_fn(cap,cur,last,fps,pl_tog):
     ret, frame_bgr = player.read(cur)
     frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
     output = segmentor(frame)
+    for i in range(len(output)):
+        mask = np.array(output[i]['mask'])/255
+        img = np.array(image)
+        img2 = np.zeros_like(img)
+        img2[:,:,0] = mask
+        img2[:,:,1] = mask
+        img2[:,:,2] = mask
+        seg_mix=img/img2
+
+        #seg_out=img2.astype(np.unit8) 
+        results=seg_mix.astype(np.uint8)
+
+
 
+
+
+
+
+    
     if pl_tog ==1:
         cur+=1
     else:
         cur = cur        
-    return output,cur
+    return results,cur
 
 with gr.Blocks() as app:
-    gr.Markdown("""<center><h1>Testing</h1><h4>stuff</h4><h4></h4></center>""")
+    gr.Markdown("""<center><h1>Slow Video Object Detection</h1><h4>Gradio and ultralyticsplus/yolov8s</h4><h4>Probably faster on GPU 🤷‍♂️</h4></center>""")
     play_state = gr.Markdown("""<right><h7></h7></right>""")
 
     with gr.Row():