Update app.py

app.py

@@ -1,15 +1,108 @@
-import spaces
-import supervision as sv
-import PIL.Image as Image
+# Import libraries
+import cv2
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download  # For Hugging Face model download
 import gradio as gr
 
-global repo_id
+# Define constants for ASL letters with color for bounding boxes
+ASL_COLORS = {
+    0: (191, 100, 21),   # A
+    1: (2, 62, 115),     # B
+    2: (140, 80, 58),    # C
+    3: (168, 181, 69),   # D
+    4: (2, 69, 84),      # E
+    5: (83, 115, 106),   # F
+    6: (255, 72, 88),    # G
+    7: (0, 204, 192),    # H
+    8: (116, 127, 127),  # I
+    9: (0, 153, 221),    # J
+    10: (196, 51, 2),    # K
+    11: (191, 100, 21),  # L
+    12: (2, 62, 115),    # M
+    13: (140, 80, 58),   # N
+    14: (168, 181, 69),  # O
+    15: (2, 69, 84),     # P
+    16: (83, 115, 106),  # Q
+    17: (255, 72, 88),   # R
+    18: (0, 204, 192),   # S
+    19: (116, 127, 127), # T
+    20: (0, 153, 221),   # U
+    21: (196, 51, 2),    # V
+    22: (191, 100, 21),  # W
+    23: (2, 62, 115),    # X
+    24: (140, 80, 58),   # Y
+    25: (168, 181, 69)   # Z
+}
+BOX_PADDING = 2
 
-repo_id = "atalaydenknalbant/asl-yolo-models"
+# Define the function to download the models dynamically
+def download_model(model_id):
+    # Use Hugging Face's hf_hub_download to download models
+    model_path = hf_hub_download(repo_id="atalaydenknalbant/asl-yolo-models", filename=model_id, local_dir="./")
+    return model_path
 
-# Model filenames directly provided, since they are known
+# Function for detecting objects (ASL letters) in the image
+def detect(image_path, model_id):
+    """
+    Output inference image with bounding boxes and ASL letter predictions.
+    Args:
+    - image_path: Path to the image file
+    - model_id: Model filename to download
+    Return: image with bounding boxes and labels drawn 
+    """
+    # Download and load the model dynamically
+    model_path = download_model(model_id)
+    detection_model = YOLO(model_path)
+
+    image = cv2.imread(image_path)
+    if image is None:
+        return image
+    
+    # Predict on image
+    results = detection_model.predict(source=image, conf=0.2, iou=0.8)  # Predict on image
+    boxes = results[0].boxes  # Get bounding boxes
+
+    if len(boxes) == 0:
+        return image
+
+    # Draw bounding boxes and labels
+    for box in boxes:
+        detection_class_conf = round(box.conf.item(), 2)  # Confidence score
+        class_id = int(box.cls.item())  # Get class ID
+        
+        # Get start and end points of the bounding box
+        start_box = (int(box.xyxy[0][0]), int(box.xyxy[0][1]))
+        end_box = (int(box.xyxy[0][2]), int(box.xyxy[0][3]))
+        
+        # 01. DRAW BOUNDING BOX OF OBJECT
+        line_thickness = round(0.001 * (image.shape[0] + image.shape[1]) / 2) + 1
+        image = cv2.rectangle(img=image, 
+                              pt1=start_box, 
+                              pt2=end_box,
+                              color=ASL_COLORS[class_id], 
+                              thickness=line_thickness)  # Draw the box with predefined colors
+        
+        # 02. DRAW LABEL
+        asl_letter = chr(65 + class_id)  # Convert class ID to ASL letter
+        text = f"{asl_letter} {detection_class_conf:.2f}"  # Label text
+        font_thickness = max(line_thickness - 1, 1)
+        (font_scale_w, font_scale_h) = (line_thickness * 0.5, line_thickness * 0.5)
+        (text_w, text_h), _ = cv2.getTextSize(text=text, fontFace=2, fontScale=font_scale_w, thickness=font_thickness)
+        
+        # Draw wrapping box for text
+        image = cv2.rectangle(img=image,
+                              pt1=(start_box[0], start_box[1] - text_h - BOX_PADDING * 2),
+                              pt2=(start_box[0] + text_w + BOX_PADDING * 2, start_box[1]),
+                              color=ASL_COLORS[class_id],
+                              thickness=-1)
+        
+        # Put class name on image
+        start_text = (start_box[0] + BOX_PADDING, start_box[1] - BOX_PADDING)
+        image = cv2.putText(img=image, text=text, org=start_text, fontFace=0, color=(255, 255, 255), fontScale=font_scale_w, thickness=font_thickness)
+
+    return image
+
+# Gradio interface
 model_filenames = [
     "yolov10s.pt",
     "yolov10x.pt",
@@ -19,129 +112,10 @@ model_filenames = [
     "yolov9s.pt"
 ]
 
-def download_models(repo_id, model_id):
-    # Download the selected model
-    hf_hub_download(repo_id, filename=model_id, local_dir=f"./")
-    return f"./{model_id}"
-
-box_annotator = sv.BoxAnnotator()
-category_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H', 8: 'I',
-                 9: 'J', 10: 'K', 11: 'L', 12: 'M', 13: 'N', 14: 'O', 15: 'P', 16: 'Q',
-                 17: 'R', 18: 'S', 19: 'T', 20: 'U', 21: 'V', 22: 'W', 23: 'X', 24: 'Y', 25: 'Z'}
-
-@spaces.GPU
-def yolo_inference(image, model_id, conf_threshold, iou_threshold, max_detection):
-    # Download models
-    model_path = download_models(repo_id, model_id)
-    
-    model = YOLO(model_path)
-    results = model(source=image, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
-    detections = sv.Detections.from_ultralytics(results)
-    
-    labels = [
-        f"{category_dict[class_id]} {confidence:.2f}"
-        for class_id, confidence in zip(detections.class_id, detections.confidence)
-    ]
-    annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)
-
-    return annotated_image
-
-def app():
-    with gr.Blocks():
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image(type="pil", label="Image", interactive=True)
-    
-                model_id = gr.Dropdown(
-                    label="Model",
-                    choices=model_filenames,
-                    value=model_filenames[0] if model_filenames else "",
-                )
-                conf_threshold = gr.Slider(
-                    label="Confidence Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.25,
-                )
-                iou_threshold = gr.Slider(
-                    label="IoU Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.45,
-                )
-                
-                max_detection = gr.Slider(
-                    label="Max Detection",
-                    minimum=1,
-                    step=1,
-                    value=1,
-                )
-                yolov_infer = gr.Button(value="Detect Objects")
-
-            with gr.Column():
-                output_image = gr.Image(type="pil", label="Annotated Image", interactive=False)
-
-        yolov_infer.click(
-            fn=yolo_inference,
-            inputs=[
-                image,
-                model_id,
-                conf_threshold,
-                iou_threshold,
-                max_detection,
-            ],
-            outputs=[output_image],
-        )
-
-        gr.Examples(
-            examples=[
-                [
-                    "b.jpg",
-                    "yolov10x.pt",
-                    0.25,
-                    0.45,
-                    1,
-                ],
-                
-                [
-                    "a.jpg",
-                    "yolov10s.pt",
-                    0.25,
-                    0.45,
-                    1,
-                ],
-                [
-                    "y.jpg",
-                    "yolov10x.pt",
-                    0.25,
-                    0.45,
-                    1,
-                ],
-            ],
-            fn=yolo_inference,
-            inputs=[
-                image,
-                model_id,
-                conf_threshold,
-                iou_threshold,
-                max_detection,
-            ],
-            outputs=[output_image],
-            cache_examples="lazy",
-        )
-
-gradio_app = gr.Blocks()
-with gradio_app:
-    gr.HTML(
-        """
-    <h1 style='text-align: center'>
-    YOLO Powered ASL(American Sign Language) Letter Detector PSA: It can't detect J or Z
-    </h1>
-        """)
-    with gr.Row():
-        with gr.Column():
-            app()
+iface = gr.Interface(fn=detect, 
+                     inputs=[gr.Image(label="Upload ASL letter image", type="filepath"), 
+                             gr.Dropdown(label="Model", choices=model_filenames, value=model_filenames[0])],
+                     outputs="image")
 
-gradio_app.launch()
+# Launch the interface
+iface.launch()