Update services/stream_object_detection_service.py

services/stream_object_detection_service.py

@@ -1,19 +1,24 @@
-from PIL import ImageDraw, ImageFont
-import spaces
+from PIL import ImageDraw, ImageFont, Image
 import cv2
-from PIL import Image
-import numpy as np
 import torch
+import numpy as np
 import uuid
+import spaces
+from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
+
+# === Load model (chỉ load 1 lần khi khởi động Space) ===
+image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd")
+model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd").to("cuda" if torch.cuda.is_available() else "cpu")
 
-SUBSAMPLE = 2  # giảm tốc độ khung hình để tăng tốc độ xử lý
+SUBSAMPLE = 2  # giảm FPS để tiết kiệm tài nguyên
 
 
-class StreamObjectDetection :
+class StreamObjectDetection:
+    @staticmethod
     def draw_bounding_boxes(image, boxes, model, conf_threshold):
         draw = ImageDraw.Draw(image)
         font = ImageFont.load_default()
-    
+
         for score, label, box in zip(boxes["scores"], boxes["labels"], boxes["boxes"]):
             if score < conf_threshold:
                 continue
@@ -21,52 +26,58 @@ class StreamObjectDetection :
             label_text = f"{model.config.id2label[label.item()]}: {score:.2f}"
             draw.rectangle([x0, y0, x1, y1], outline="red", width=3)
             draw.text((x0 + 3, y0 + 3), label_text, fill="white", font=font)
-    
+
         return image
-        
-    @spaces.CPU  # dùng GPU nếu chạy trên Hugging Face
-    def stream_object_detection(video, conf_threshold):
+
+    @staticmethod
+    @spaces.GPU  # Dùng GPU nếu có (ZeroGPU, GPU Cluster, v.v.)
+    def stream_object_detection(video, conf_threshold=0.3):
         cap = cv2.VideoCapture(video)
         video_codec = cv2.VideoWriter_fourcc(*"mp4v")
-        fps = int(cap.get(cv2.CAP_PROP_FPS))
-        desired_fps = fps // SUBSAMPLE
+        fps = int(cap.get(cv2.CAP_PROP_FPS)) or 24
+        desired_fps = max(1, fps // SUBSAMPLE)
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2
-    
+
         iterating, frame = cap.read()
         n_frames = 0
         output_video_name = f"output_{uuid.uuid4()}.mp4"
         output_video = cv2.VideoWriter(output_video_name, video_codec, desired_fps, (width, height))
         batch = []
-    
+
         while iterating:
             frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    
+
             if n_frames % SUBSAMPLE == 0:
                 batch.append(frame)
-    
-            if len(batch) == 2 * desired_fps:  # mỗi 2s xử lý 1 lần
+
+            # Mỗi 2 giây xử lý một lần
+            if len(batch) == 2 * desired_fps:
                 inputs = image_processor(images=batch, return_tensors="pt").to(model.device)
+
                 with torch.no_grad():
                     outputs = model(**inputs)
-    
+
                 boxes = image_processor.post_process_object_detection(
                     outputs,
                     target_sizes=torch.tensor([(height, width)] * len(batch)).to(model.device),
                     threshold=conf_threshold,
                 )
-    
+
                 for img, box in zip(batch, boxes):
-                    pil_image = draw_bounding_boxes(Image.fromarray(img), box, model, conf_threshold)
-                    frame = np.array(pil_image)[:, :, ::-1]
-                    output_video.write(frame)
-    
+                    pil_image = StreamObjectDetection.draw_bounding_boxes(Image.fromarray(img), box, model, conf_threshold)
+                    frame_bgr = np.array(pil_image)[:, :, ::-1]
+                    output_video.write(frame_bgr)
+
                 batch = []
                 output_video.release()
-                yield output_video_name  # stream ra video đã xử lý
+                yield output_video_name  # Gửi video xử lý từng phần cho Gradio
                 output_video_name = f"output_{uuid.uuid4()}.mp4"
                 output_video = cv2.VideoWriter(output_video_name, video_codec, desired_fps, (width, height))
-    
+
             iterating, frame = cap.read()
-            n_frames += 1
+            n_frames += 1
+
+        cap.release()
+        output_video.release()