Update app.py

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 # Muat model pre-trained YOLOv5
 model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 
-# Fungsi untuk memproses video dan menghitung jumlah manusia
+# Fungsi untuk memproses video
 def process_video(video_path):
     # Direktori output
     output_dir = "output_videos"
@@ -26,7 +26,6 @@ def process_video(video_path):
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
     out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
 
-    # Proses video
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
@@ -56,21 +55,55 @@ def process_video(video_path):
     cap.release()
     out.release()
 
-    # Mengembalikan video yang telah diproses (tidak menjumlahkan seluruh frame)
     return output_path
 
-# Fungsi Gradio untuk antarmuka
-def gradio_interface(video_file):
-    output_path = process_video(video_file)
+# Fungsi untuk memproses gambar
+def process_image(image_path):
+    # Baca gambar
+    image = cv2.imread(image_path)
+
+    # Inferensi dengan YOLOv5
+    results = model(image)
+    detections = results.pred[0]
+    names = model.names
+
+    # Filter hanya label 'person'
+    person_detections = [d for d in detections if names[int(d[-1])] == "person"]
+    person_count = len(person_detections)
+
+    # Render frame dan buat salinan eksplisit
+    annotated_image = results.render()[0]
+    annotated_image = np.copy(annotated_image)
+
+    # Tambahkan teks ke gambar
+    cv2.putText(annotated_image, f"Person Count: {person_count}", (10, 30),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+    # Simpan gambar hasil
+    output_dir = "output_images"
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = os.path.join(output_dir, "person_counter_output.jpg")
+    cv2.imwrite(output_path, annotated_image)
+
     return output_path
 
+# Fungsi Gradio untuk antarmuka
+def gradio_interface(file, is_video):
+    if is_video:
+        return process_video(file)
+    else:
+        return process_image(file)
+
 # Antarmuka Gradio
 iface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.File(type="filepath"),  # Input berupa file video
-    outputs=gr.File(label="Processed Video"),
+    inputs=[
+        gr.File(type="filepath", label="Upload File (Image/Video)"),
+        gr.Checkbox(label="Is Video?", value=True),
+    ],
+    outputs=gr.File(label="Processed File"),
     title="Person Counter using YOLOv5",
-    description="Upload a video file to detect and count the number of people in each frame using YOLOv5."
+    description="Upload a video or image file to detect and count the number of people using YOLOv5."
 )
 
 # Menjalankan aplikasi