Update app.py

app.py

@@ -16,102 +16,124 @@ Show your appreciation for this space-age tool by hitting the 'Like' button and
 📧 Contact us: info@foddu.com
 👍 Like | """
 
-def download_file(url, save_name):
-    url = url
-    if not os.path.exists(save_name):
-        file = requests.get(url)
-        open(save_name, 'wb').write(file.content)
-
-# Download files
-file_urls = [
-    'https://huggingface.co/spaces/foduucom/CandleStickScan-Stock-trading-yolov8/resolve/main/test/-2022-06-28-12-35-50_png.rf.8dee4bb645ea8b5036721b830d2636b1.jpg',
-    'https://huggingface.co/spaces/foduucom/CandleStickScan-Stock-trading-yolov8/resolve/main/test/-2022-06-28-12-45-10_png.rf.8b9177546e62a2422ad603b16f1f50b9.jpg',
-    'https://www.dropbox.com/s/7sjfwncffg8xej2/video_7.mp4?dl=1'
-]
-
-for i, url in enumerate(file_urls):
-    if 'mp4' in file_urls[i]:
-        download_file(
-            file_urls[i],
-            f"video.mp4"
-        )
-    # else:
-    #     download_file(
-    #         file_urls[i],
-    #         f"image_{i}.jpg"
-    #     )
+image_path= [['test/test1.jpg', 'foduucom/stockmarket-pattern-detection-yolov8', 640, 0.25, 0.45], ['test/test2.jpg', 'foduucom/stockmarket-pattern-detection-yolov8', 640, 0.25, 0.45]]
 
 # Load YOLO model
 model = YOLO('foduucom/stockmarket-pattern-detection-yolov8')
 
-# Image Inference
-def show_preds_image(image_path):
-    image = cv2.imread(image_path)
-    outputs = model.predict(source=image_path)
-    results = outputs[0].cpu().numpy()
-    for i, det in enumerate(results.boxes.xyxy):
-        cv2.rectangle(
-            image,
-            (int(det[0]), int(det[1])),
-            (int(det[2]), int(det[3])),
-            color=(0, 0, 255),
-            thickness=2,
-            lineType=cv2.LINE_AA
-        )
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+#############################################################Image Inference############################################################
+def yolov8_img_inference(
+    image: gr.inputs.Image = None,
+    model_path: gr.inputs.Dropdown = None,
+    image_size: gr.inputs.Slider = 640,
+    conf_threshold: gr.inputs.Slider = 0.25,
+    iou_threshold: gr.inputs.Slider = 0.45,
+):
+    """
+    YOLOv8 inference function
+    Args:
+        image: Input image
+        model_path: Path to the model
+        image_size: Image size
+        conf_threshold: Confidence threshold
+        iou_threshold: IOU threshold
+    Returns:
+        Rendered image
+    """
+    model = YOLO(model_path)
+    model.overrides['conf'] = conf_threshold
+    model.overrides['iou']= iou_threshold
+    model.overrides['agnostic_nms'] = False  # NMS class-agnostic
+    model.overrides['max_det'] = 1000 
+    image = read_image(image)
+    results = model.predict(image)
+    render = render_result(model=model, image=image, result=results[0])
+    
+    return render
 
+    
 inputs_image = [
-    gr.components.Image(type="filepath", label="Input Image"),
-]
-outputs_image = [
-    gr.components.Image(type="numpy", label="Output Image"),
+    gr.inputs.Image(type="filepath", label="Input Image"),
+    gr.inputs.Dropdown(["foduucom/stockmarket-pattern-detection-yolov8"], 
+                       default="foduucom/stockmarket-pattern-detection-yolov8", label="Model"),
+    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
 ]
+
+outputs_image =gr.outputs.Image(type="filepath", label="Output Image")
 interface_image = gr.Interface(
-    fn=show_preds_image,
+    fn=yolov8_img_inference,
     inputs=inputs_image,
     outputs=outputs_image,
     title=model_heading,
-    descripiton=description,
-    examples=path,
+    description=description,
+    examples=image_path,
     cache_examples=False,
 )
 
-# Video Inference
-def show_preds_video(video_path):
+##################################################Video Inference################################################################
+def show_preds_video(
+    video_path: gr.components.Video = None,
+    model_path: gr.inputs.Dropdown = None,
+    image_size: gr.inputs.Slider = 640,
+    conf_threshold: gr.inputs.Slider = 0.25,
+    iou_threshold: gr.inputs.Slider = 0.45,
+):
+    """
+    Video inference function
+    Args:
+        video_path: Input video
+        model_path: Path to the model
+        image_size: Image size
+        conf_threshold: Confidence threshold
+        iou_threshold: IOU threshold
+    Returns:
+        Rendered video
+    """
     cap = cv2.VideoCapture(video_path)
-    while(cap.isOpened()):
-        ret, frame = cap.read()
-        if ret:
-            frame_copy = frame.copy()
-            outputs = model.predict(source=frame)
-            results = outputs[0].cpu().numpy()
-            for i, det in enumerate(results.boxes.xyxy):
-                cv2.rectangle(
-                    frame_copy,
-                    (int(det[0]), int(det[1])),
-                    (int(det[2]), int(det[3])),
-                    color=(0, 0, 255),
-                    thickness=2,
-                    lineType=cv2.LINE_AA
-                )
-            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
 
+    while cap.isOpened():
+        success, frame = cap.read()
+
+        if success:
+            model = YOLO(model_path)
+            model.overrides['conf'] = conf_threshold
+            model.overrides['iou'] = iou_threshold
+            model.overrides['agnostic_nms'] = False
+            model.overrides['max_det'] = 1000 
+            results = model.predict(frame)
+            annotated_frame = results[0].plot()
+
+            cv2.imshow("YOLOv8 Inference", annotated_frame)
+
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                break
+        else:
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+    
 inputs_video = [
     gr.components.Video(type="filepath", label="Input Video"),
+    gr.inputs.Dropdown(["foduucom/stockmarket-pattern-detection-yolov8"], 
+                       default="foduucom/stockmarket-pattern-detection-yolov8", label="Model"),
+    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
 
 ]
-outputs_video = [
-    gr.components.Image(type="numpy", label="Output Image"),
-]
-
+outputs_video = gr.outputs.Image(type="filepath", label="Output Video")
+video_path=[['test/video.mp4','foduucom/stockmarket-pattern-detection-yolov8', 640, 0.25, 0.45]]
 interface_video = gr.Interface(
     fn=show_preds_video,
     inputs=inputs_video,
     outputs=outputs_video,
     title=model_heading,
-    descripiton=description,
+    description=description,
     examples=video_path,
-    cache_examples=False,
+    cache_examples=True,
 )
 
 gr.TabbedInterface(