Update app.py

app.py

@@ -4,6 +4,7 @@ from roboflow import Roboflow
 import tempfile
 import os
 import requests
+import cv2
 
 # Muat variabel lingkungan dari file .env
 load_dotenv()
@@ -17,7 +18,7 @@ rf = Roboflow(api_key=api_key)
 project = rf.workspace(workspace).project(project_name)
 model = project.version(model_version).model
 
-# Fungsi untuk menangani input dan output gambar
+# Fungsi untuk menangani deteksi pada gambar
 def detect_objects(image):
     # Simpan gambar yang diupload sebagai file sementara
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
@@ -26,16 +27,16 @@ def detect_objects(image):
 
     try:
         # Lakukan prediksi pada gambar
-        predictions = model.predict(temp_file_path, confidence=50, overlap=80).json()
+        predictions = model.predict(temp_file_path, confidence=60, overlap=80).json()
 
         # Menghitung jumlah objek per kelas
         class_count = {}
-        total_count = 0  # Menyimpan total jumlah objek
+        total_count = 0
 
         for prediction in predictions['predictions']:
             class_name = prediction['class']
             class_count[class_name] = class_count.get(class_name, 0) + 1
-            total_count += 1  # Tambah jumlah objek untuk setiap prediksi
+            total_count += 1
 
         # Menyusun output berupa string hasil perhitungan
         result_text = "Product Nestle\n\n"
@@ -45,41 +46,97 @@ def detect_objects(image):
 
         # Menyimpan gambar dengan prediksi
         output_image_path = "/tmp/prediction.jpg"
-        model.predict(temp_file_path, confidence=50, overlap=80).save(output_image_path)
+        model.predict(temp_file_path, confidence=60, overlap=80).save(output_image_path)
         
     except requests.exceptions.HTTPError as http_err:
-        # Menangani kesalahan HTTP
         result_text = f"HTTP error occurred: {http_err}"
-        output_image_path = temp_file_path  # Kembalikan gambar asli jika terjadi error
+        output_image_path = temp_file_path
     except Exception as err:
-        # Menangani kesalahan lain
         result_text = f"An error occurred: {err}"
-        output_image_path = temp_file_path  # Kembalikan gambar asli jika terjadi error
+        output_image_path = temp_file_path
 
-    # Hapus file sementara setelah prediksi
     os.remove(temp_file_path)
     
     return output_image_path, result_text
 
+# Fungsi untuk menangani deteksi pada video
+def detect_objects_in_video(video_path):
+    temp_output_path = "/tmp/output_video.mp4"
+    temp_frames_dir = tempfile.mkdtemp()
+
+    try:
+        # Baca video dan ekstrak frame
+        video = cv2.VideoCapture(video_path)
+        frame_rate = int(video.get(cv2.CAP_PROP_FPS))
+        frame_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+        frame_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        frame_size = (frame_width, frame_height)
+        frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        # VideoWriter untuk membuat video keluaran
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        output_video = cv2.VideoWriter(temp_output_path, fourcc, frame_rate, frame_size)
+
+        frame_index = 0
+        while True:
+            ret, frame = video.read()
+            if not ret:
+                break
+
+            # Simpan frame sementara untuk prediksi
+            frame_path = os.path.join(temp_frames_dir, f"frame_{frame_index}.jpg")
+            cv2.imwrite(frame_path, frame)
+
+            # Deteksi objek pada frame
+            predictions = model.predict(frame_path, confidence=60, overlap=80).json()
+
+            # Gambar bounding box pada frame
+            for prediction in predictions['predictions']:
+                x, y, w, h = prediction['x'], prediction['y'], prediction['width'], prediction['height']
+                class_name = prediction['class']
+                color = (0, 255, 0)  # Hijau
+                cv2.rectangle(frame, (int(x - w/2), int(y - h/2)), (int(x + w/2), int(y + h/2)), color, 2)
+                cv2.putText(frame, class_name, (int(x - w/2), int(y - h/2 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+            # Tambahkan frame ke video keluaran
+            output_video.write(frame)
+            frame_index += 1
+
+        video.release()
+        output_video.release()
+
+        return temp_output_path, "Video processing completed successfully."
+    
+    except Exception as e:
+        return None, f"An error occurred: {e}"
+
 # Membuat antarmuka Gradio dengan tata letak fleksibel
 with gr.Blocks() as iface:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil", label="Input Image")
+            input_video = gr.Video(type="file", label="Input Video")
         with gr.Column():
             output_image = gr.Image(label="Detect Object")
+            output_video = gr.Video(label="Output Video")
         with gr.Column():
             output_text = gr.Textbox(label="Counting Object")
     
-    # Tombol untuk memproses input
-    detect_button = gr.Button("Detect")
-    
-    # Hubungkan tombol dengan fungsi deteksi
-    detect_button.click(
+    # Tombol untuk memproses gambar
+    detect_image_button = gr.Button("Detect Image")
+    detect_image_button.click(
         fn=detect_objects, 
         inputs=input_image, 
         outputs=[output_image, output_text]
     )
 
+    # Tombol untuk memproses video
+    detect_video_button = gr.Button("Detect Video")
+    detect_video_button.click(
+        fn=detect_objects_in_video, 
+        inputs=input_video, 
+        outputs=[output_video, output_text]
+    )
+
 # Menjalankan antarmuka
-iface.launch()
+iface.launch()