app.py

import gradio as gr
from ultralytics import YOLO
import os
from datetime import datetime
import requests
from requests.auth import HTTPBasicAuth
import base64
import cv2
from dotenv import load_dotenv
import io

model = YOLO('head_yolov8s.pt')
load_dotenv()

def upload_to_github(file_content, filepath):
    username = "mohitsriv99"
    repository = "hsamples"
    branch = "main"
    url = f"https://api.github.com/repos/{username}/{repository}/contents/{filepath}"
    token = os.getenv("TOKEN")  # Using an environment variable for the token

    encoded_file = base64.b64encode(file_content).decode("utf-8")

    payload = {
        "message": f"Upload image {filepath}",
        "content": encoded_file,
        "branch": branch
    }

    headers = {
        "Authorization": f"token {token}"
    }

    response = requests.put(url, json=payload, headers=headers)
    if response.status_code == 201:
        print("File uploaded successfully to GitHub.")
    else:
        print(f"Failed to upload file. Status code: {response.status_code}, Response: {response.text}")

def head_detect(path):
    if path is None:
        return None
    
    img = cv2.imread(path)

    # Generate a timestamped filename
    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"uploaded_images/{current_time}.jpg"

    # Encode the image to a memory buffer instead of saving to disk
    _, buffer = cv2.imencode('.jpg', img)
    img_byte_array = io.BytesIO(buffer).getvalue()

    # Upload the image to GitHub
    upload_to_github(img_byte_array, filename)

    output = model(source=img)
    res = output[0].cpu().numpy()
    
    # Extract bbox, cls id and conf
    bboxes = res.boxes.xyxy
    class_ids = res.boxes.cls
    conf_scores = res.boxes.conf

    for i in range(len(bboxes)):
        xmin, ymin, xmax, ymax = int(bboxes[i][0]), int(bboxes[i][1]), int(bboxes[i][2]), int(bboxes[i][3])
        conf = conf_scores[i]
        cls_id = int(class_ids[i])
        label = model.names[cls_id]  # Get the label name

        # Draw rectangle for bounding box
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)

        # Prepare label text with confidence score
        label_text = f'{label} {conf:.2f}'

        # Put text (label) on the image
        cv2.putText(img, label_text, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, lineType=cv2.LINE_AA)

    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

input_image = [
    gr.components.Image(type='filepath', label='Input Image'),
]

output_image = [
    gr.components.Image(type='numpy', label='Prediction'),
]

interface = gr.Interface(
    fn=head_detect,
    inputs=input_image,
    outputs=output_image,
    title='Top Head Detection',
    cache_examples=False,
)

gr.TabbedInterface(
    [interface],
    tab_names=['Image Inference']
).queue().launch()